├── .github └── workflows │ └── build_and_test.yml ├── .gitignore ├── LICENSE ├── Makefile ├── NOTICE ├── README.md ├── distfunc.c ├── embedding--0.3.5--0.3.6.sql ├── embedding--0.3.5.sql ├── embedding--0.3.6.sql ├── embedding.c ├── embedding.control ├── embedding.h ├── hnswalg.cpp ├── hnswalg.h └── test ├── expected ├── gh-2.out ├── gh-3.out └── knn.out └── sql ├── gh-2.sql ├── gh-3.sql └── knn.sql /.github/workflows/build_and_test.yml: -------------------------------------------------------------------------------- 1 | name: Build and Test 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | 9 | concurrency: 10 | group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }} 11 | cancel-in-progress: true 12 | 13 | jobs: 14 | build-and-test: 15 | strategy: 16 | fail-fast: false 17 | matrix: 18 | os: [ ubuntu-22.04 ] 19 | postgresql: [ 12, 13, 14, 15, 16 ] 20 | include: 21 | - os: macos-13 22 | postgresql: 15 23 | - os: macos-13 24 | postgresql: 16 25 | 26 | runs-on: ${{ matrix.os }} 27 | 28 | steps: 29 | - uses: Homebrew/actions/setup-homebrew@master 30 | 31 | - name: Workaround GitHub Actions Python issues 32 | if: startsWith(matrix.os, 'macos-') 33 | run: brew unlink python && brew link --overwrite python 34 | 35 | # TODO: Remove after postgresql@16 is available in homebrew/core 36 | - if: matrix.postgresql == 16 37 | run: brew tap bayandin/tap 38 | 39 | - name: brew install postgresql 40 | run: | 41 | brew install ${FORMULA} 42 | brew services start ${FORMULA} 43 | 44 | $(brew --prefix ${FORMULA})/bin/postgres --version 45 | 46 | echo "PG_CONFIG=$(brew --prefix ${FORMULA})/bin/pg_config" >> $GITHUB_ENV 47 | env: 48 | FORMULA: postgresql@${{ matrix.postgresql }} 49 | 50 | - uses: actions/checkout@v3 51 | 52 | - run: make 53 | - run: make install 54 | - run: make installcheck 55 | 56 | - if: failure() 57 | run: cat regression.diffs || true 58 | 59 | build-and-test-on-neon: 60 | strategy: 61 | fail-fast: false 62 | matrix: 63 | os: [ ubuntu-22.04 ] 64 | postgresql: [ 14, 15 ] 65 | include: 66 | - os: macos-13 67 | postgresql: 15 68 | 69 | runs-on: ${{ matrix.os }} 70 | 71 | steps: 72 | - uses: Homebrew/actions/setup-homebrew@master 73 | 74 | - name: Workaround GitHub Actions Python issues 75 | if: startsWith(matrix.os, 'macos-') 76 | run: brew unlink python && brew link --overwrite python 77 | 78 | - run: brew tap bayandin/tap 79 | 80 | - name: brew install neon 81 | run: | 82 | brew install ${FORMULA} 83 | 84 | echo "$(brew --prefix openssl@3)/bin" >> $GITHUB_PATH 85 | echo "PG_CONFIG=$(brew ruby -e 'puts "neon-postgres".to_sym.f.pg_bin_for("v${{ matrix.postgresql }}")')/pg_config" >> $GITHUB_ENV 86 | 87 | $(brew --prefix ${FORMULA})/bin/neon-local --version 88 | env: 89 | FORMULA: bayandin/tap/neon-local 90 | 91 | - name: Start neon 92 | run: | 93 | neon-local start 94 | neon-local tenant create --pg-version ${{ matrix.postgresql }} --set-default 95 | neon-local endpoint start main --pg-version ${{ matrix.postgresql }} --pg-port 60000 96 | 97 | - uses: actions/checkout@v3 98 | 99 | - run: make 100 | - run: make install 101 | - name: make installcheck 102 | run: | 103 | export PGPORT=60000 104 | export PGUSER=cloud_admin 105 | 106 | make installcheck 107 | 108 | - if: failure() 109 | run: cat regression.diffs || true 110 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # `make` artifacts 2 | *.so 3 | *.o 4 | *.dylib 5 | *.bc 6 | .deps/ 7 | 8 | # `make installcheck` artifacts 9 | results/ 10 | regression.diffs 11 | regression.out 12 | 13 | # `make dist` artifacts 14 | dist/ 15 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | EXTENSION = embedding 2 | EXTVERSION = 0.3.6 3 | 4 | MODULE_big = embedding 5 | DATA = $(wildcard *--*.sql) 6 | OBJS = embedding.o hnswalg.o distfunc.o 7 | 8 | TESTS = $(wildcard test/sql/*.sql) 9 | REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS)) 10 | REGRESS_OPTS = --inputdir=test --load-extension=embedding 11 | 12 | # For auto-vectorization: 13 | # - GCC&clang needs -Ofast or -O3: https://gcc.gnu.org/projects/tree-ssa/vectorization.html 14 | PG_CFLAGS += -Ofast 15 | PG_CXXFLAGS += -std=c++11 16 | PG_LDFLAGS += -lstdc++ 17 | 18 | all: $(EXTENSION)--$(EXTVERSION).sql 19 | 20 | PG_CONFIG ?= pg_config 21 | PGXS := $(shell $(PG_CONFIG) --pgxs) 22 | include $(PGXS) 23 | 24 | dist: 25 | mkdir -p dist 26 | git archive --format zip --prefix=$(EXTENSION)-$(EXTVERSION)/ --output dist/$(EXTENSION)-$(EXTVERSION).zip main 27 | -------------------------------------------------------------------------------- /NOTICE: -------------------------------------------------------------------------------- 1 | pg_embedding 2 | Copyright 2023 Neon Inc. 3 | 4 | This product includes software developed at 5 | Neon Inc. (https://www.neon.tech/). 6 | 7 | ------------------------------------------------------------------------------- 8 | HNSW Algorithm implementation (hnswalg.cpp and hnswalg.h) is derived from https://github.com/dbaranchuk/ivf-hnsw, 9 | which is licensed under the MIT License. 10 | 11 | MIT License 12 | 13 | Copyright (c) 2017 Dmitry Baranchuk 14 | 15 | Permission is hereby granted, free of charge, to any person obtaining a copy 16 | of this software and associated documentation files (the "Software"), to deal 17 | in the Software without restriction, including without limitation the rights 18 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 19 | copies of the Software, and to permit persons to whom the Software is 20 | furnished to do so, subject to the following conditions: 21 | 22 | The above copyright notice and this permission notice shall be included in all 23 | copies or substantial portions of the Software. 24 | 25 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 26 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 27 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 28 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 29 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 30 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 31 | SOFTWARE. 32 | ------------------------------------------------------------------------------- 33 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pg_embedding 2 | 3 | --- 4 | **IMPORTANT NOTICE:** 5 | 6 | As of Sept 29, 2023, Neon is no longer committing to `pg_embedding`. 7 | 8 | Support will remain in place for existing users of the extension, but we strongly encourage migrating to `pgvector`. 9 | 10 | For migration instructions, see [Migrate from pg_embedding to pgvector](https://neon.tech/docs/extensions/pg_embedding#migrate-from-pg_embedding-to-pgvector), in the _Neon documentation_. 11 | 12 | --- 13 | 14 | The `pg_embedding` extension enables the using the Hierarchical Navigable Small World (HNSW) algorithm for vector similarity search in PostgreSQL. 15 | 16 | This extension is based on [ivf-hnsw](https://github.com/dbaranchuk/ivf-hnsw) implementation of HNSW 17 | the code for the current state-of-the-art billion-scale nearest neighbor search system[[1]](#references). 18 | 19 | ## Using the pg_embedding extension 20 | 21 | This section describes how to use the `pg_embedding` extension with a simple example demonstrating the required statements, syntax, and options. 22 | 23 | For information about migrating from `pgvector` to `pg_embedding`, see [Migrate from pgvector to pg_embedding](https://neon.tech/docs/extensions/pg_embedding#migrate-from-pgvector-to-pgembedding), in the _Neon documentation_. 24 | 25 | ### Usage summary 26 | 27 | The statements in this usage summary are described in further detail in the following sections. 28 | 29 | ```sql 30 | CREATE EXTENSION embedding; 31 | CREATE TABLE documents(id integer PRIMARY KEY, embedding real[]); 32 | INSERT INTO documents(id, embedding) VALUES (1, '{0,1,2}'), (2, '{1,2,3}'), (3, '{1,1,1}'); 33 | SELECT id FROM documents ORDER BY embedding <-> ARRAY[3,3,3] LIMIT 1; 34 | ``` 35 | 36 | ### Enable the extension 37 | 38 | To enable the `pg_embedding` extension, run the following `CREATE EXTENSION` statement: 39 | 40 | ```sql 41 | CREATE EXTENSION embedding; 42 | ``` 43 | 44 | ### Create a table for your vector data 45 | 46 | To store your vector data, create a table similar to the following: 47 | 48 | ```sql 49 | CREATE TABLE documents(id INTEGER, embedding REAL[]); 50 | ``` 51 | 52 | This statement generates a table named `documents` with an `embedding` column for storing vector data. Your table and vector column names may differ. 53 | 54 | ### Insert data 55 | 56 | To insert vector data, use an `INSERT` statement similar to the following: 57 | 58 | ```sql 59 | INSERT INTO documents(id, embedding) VALUES (1, '{0,1,2}'), (2, '{1,2,3}'), (3, '{1,1,1}'); 60 | ``` 61 | 62 | ## Query 63 | 64 | The `pg_embedding` extension supports Euclidean (L2), Cosine, and Manhattan distance metrics. 65 | 66 | Euclidean (L2) distance: 67 | 68 | ```sql 69 | SELECT id FROM documents ORDER BY embedding <-> array[3,3,3] LIMIT 1; 70 | ``` 71 | 72 | Cosine distance: 73 | 74 | ```sql 75 | SELECT id FROM documents ORDER BY embedding <=> array[3,3,3] LIMIT 1; 76 | ``` 77 | 78 | Manhattan distance: 79 | 80 | ```sql 81 | SELECT id FROM documents ORDER BY embedding <~> array[3,3,3] LIMIT 1; 82 | ``` 83 | 84 | where: 85 | 86 | - `SELECT id FROM documents` selects the `id` field from all records in the `documents` table. 87 | - `ORDER BY` sorts the selected records in ascending order based on the calculated distances. In other words, records with values closer to the `[1.1, 2.2, 3.3]` query vector according to the distance metric will be returned first. 88 | - `<->`, `<=>`, and `<~>` operators define the distance metric, which calculates the distance between the query vector and each row of the dataset. 89 | - `LIMIT 1` limits the result set to one record after sorting. 90 | 91 | In summary, the query retrieves the ID of the record from the `documents` table whose value is closest to the `[3,3,3]` query vector according to the specified distance metric. 92 | 93 | ### Create an HNSW index 94 | 95 | To optimize search behavior, you can add an HNSW index. To create the HNSW index on your vector column, use a `CREATE INDEX` statement as shown in the following examples. The `pg_embedding` extension supports indexes for use with Euclidean, Cosine, and Manhattan distance metrics. 96 | 97 | Euclidean (L2) distance index: 98 | 99 | ```sql 100 | CREATE INDEX ON documents USING hnsw(embedding) WITH (dims=3, m=3, efconstruction=5, efsearch=5); 101 | SET enable_seqscan = off; 102 | SELECT id FROM documents ORDER BY embedding <-> array[3,3,3] LIMIT 1; 103 | ``` 104 | 105 | Cosine distance index: 106 | 107 | ```sql 108 | CREATE INDEX ON documents USING hnsw(embedding ann_cos_ops) WITH (dims=3, m=3, efconstruction=5, efsearch=5); 109 | SET enable_seqscan = off; 110 | SELECT id FROM documents ORDER BY embedding <=> array[3,3,3] LIMIT 1; 111 | ``` 112 | 113 | Manhattan distance index: 114 | 115 | ```sql 116 | CREATE INDEX ON documents USING hnsw(embedding ann_manhattan_ops) WITH (dims=3, m=3, efconstruction=5, efsearch=5); 117 | SET enable_seqscan = off; 118 | SELECT id FROM documents ORDER BY embedding <~> array[3,3,3] LIMIT 1; 119 | ``` 120 | 121 | ### Tuning the HNSW algorithm 122 | 123 | The following options allow you to tune the HNSW algorithm when creating an index: 124 | 125 | - `dims`: Defines the number of dimensions in your vector data. This is a required parameter. 126 | - `m`: Defines the maximum number of links or "edges" created for each node during graph construction. A higher value increases accuracy (recall) but also increases the size of the index in memory and index construction time. 127 | - `efconstruction`: Influences the trade-off between index quality and construction speed. A high `efconstruction` value creates a higher quality graph, enabling more accurate search results, but a higher value also means that index construction takes longer. 128 | - `efsearch`: Influences the trade-off between query accuracy (recall) and speed. A higher `efsearch` value increases accuracy at the cost of speed. This value should be equal to or larger than `k`, which is the number of nearest neighbors you want your search to return (defined by the `LIMIT` clause in your `SELECT` query). 129 | 130 | In summary, to prioritize search speed over accuracy, use lower values for `m` and `efsearch`. Conversely, to prioritize accuracy over search speed, use a higher value for `m` and `efsearch`. A higher `efconstruction` value enables more accurate search results at the cost of index build time, which is also affected by the size of your dataset. 131 | 132 | ## How HNSW search works 133 | 134 | HNSW is a graph-based approach to indexing multi-dimensional data. It constructs a multi-layered graph, where each layer is a subset of the previous one. During a search, the algorithm navigates through the graph from the top layer to the bottom to quickly find the nearest neighbor. An HNSW graph is known for its superior performance in terms of speed and accuracy. 135 | 136 | The search process begins at the topmost layer of the HNSW graph. From the starting node, the algorithm navigates to the nearest neighbor in the same layer. The algorithm repeats this step until it can no longer find neighbors more similar to the query vector. 137 | 138 | Using the found node as an entry point, the algorithm moves down to the next layer in the graph and repeats the process of navigating to the nearest neighbor. The process of navigating to the nearest neighbor and moving down a layer is repeated until the algorithm reaches the bottom layer. 139 | 140 | In the bottom layer, the algorithm continues navigating to the nearest neighbor until it can't find any nodes that are more similar to the query vector. The current node is then returned as the most similar node to the query vector. 141 | 142 | The key idea behind HNSW is that by starting the search at the top layer and moving down through each layer, the algorithm can quickly navigate to the area of the graph that contains the node that is most similar to the query vector. This makes the search process much faster than if it had to search through every node in the graph. 143 | 144 | ## References 145 | 146 | - [1] Dmitry Baranchuk, Artem Babenko, Yury Malkov; Proceedings of the European Conference on Computer Vision (ECCV), 2018, pp. 202-216 [link](http://openaccess.thecvf.com/content_ECCV_2018/html/Dmitry_Baranchuk_Revisiting_the_Inverted_ECCV_2018_paper.html) 147 | -------------------------------------------------------------------------------- /distfunc.c: -------------------------------------------------------------------------------- 1 | // Copyright 2023 Neon Inc. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #include "postgres.h" 16 | #include "embedding.h" 17 | #include "math.h" 18 | 19 | #ifdef __x86_64__ 20 | #include 21 | 22 | #if defined(__GNUC__) 23 | #define PORTABLE_ALIGN32 __attribute__((aligned(32))) 24 | #else 25 | #define PORTABLE_ALIGN32 __declspec(align(32)) 26 | #endif 27 | 28 | __attribute__((target("avx2"))) 29 | static dist_t l2_dist_impl_avx2(const coord_t *x, const coord_t *y, size_t n) 30 | { 31 | coord_t PORTABLE_ALIGN32 TmpRes[sizeof(__m256) / sizeof(float)]; 32 | size_t qty16 = n / 16; 33 | const coord_t *pEnd1 = x + (qty16 * 16); 34 | const coord_t *pEnd2 = x + n; 35 | __m256 diff, v1, v2; 36 | __m256 sum = _mm256_set1_ps(0); 37 | dist_t res; 38 | 39 | while (x < pEnd1) { 40 | v1 = _mm256_loadu_ps(x); 41 | x += 8; 42 | v2 = _mm256_loadu_ps(y); 43 | y += 8; 44 | diff = _mm256_sub_ps(v1, v2); 45 | sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff)); 46 | 47 | v1 = _mm256_loadu_ps(x); 48 | x += 8; 49 | v2 = _mm256_loadu_ps(y); 50 | y += 8; 51 | diff = _mm256_sub_ps(v1, v2); 52 | sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff)); 53 | } 54 | _mm256_store_ps(TmpRes, sum); 55 | res = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + TmpRes[5] + TmpRes[6] + TmpRes[7]; 56 | 57 | // Handle case when dimensions is not aligned on 16. 58 | while (x < pEnd2) 59 | { 60 | dist_t diff = *x++ - *y++; 61 | res += diff * diff; 62 | } 63 | 64 | return sqrtf(res); 65 | } 66 | 67 | static dist_t l2_dist_impl_sse(const coord_t *x, const coord_t *y, size_t n) 68 | { 69 | coord_t PORTABLE_ALIGN32 TmpRes[sizeof(__m128) / sizeof(float)]; 70 | size_t qty16 = n / 16; 71 | const coord_t *pEnd1 = x + (qty16 * 16); 72 | const coord_t *pEnd2 = x + n; 73 | dist_t res; 74 | 75 | __m128 diff, v1, v2; 76 | __m128 sum = _mm_set1_ps(0); 77 | 78 | while (x < pEnd1) { 79 | v1 = _mm_loadu_ps(x); 80 | x += 4; 81 | v2 = _mm_loadu_ps(y); 82 | y += 4; 83 | diff = _mm_sub_ps(v1, v2); 84 | sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); 85 | 86 | v1 = _mm_loadu_ps(x); 87 | x += 4; 88 | v2 = _mm_loadu_ps(y); 89 | y += 4; 90 | diff = _mm_sub_ps(v1, v2); 91 | sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); 92 | 93 | v1 = _mm_loadu_ps(x); 94 | x += 4; 95 | v2 = _mm_loadu_ps(y); 96 | y += 4; 97 | diff = _mm_sub_ps(v1, v2); 98 | sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); 99 | 100 | v1 = _mm_loadu_ps(x); 101 | x += 4; 102 | v2 = _mm_loadu_ps(y); 103 | y += 4; 104 | diff = _mm_sub_ps(v1, v2); 105 | sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); 106 | } 107 | _mm_store_ps(TmpRes, sum); 108 | res = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3]; 109 | 110 | // Handle case when dimensions is not aligned on 16. 111 | while (x < pEnd2) 112 | { 113 | dist_t diff = *x++ - *y++; 114 | res += diff * diff; 115 | } 116 | 117 | return sqrtf(res); 118 | } 119 | #else 120 | 121 | static dist_t l2_dist_impl(coord_t const* ax, coord_t const* bx, size_t dim) 122 | { 123 | dist_t distance = 0.0; 124 | for (size_t i = 0; i < dim; i++) 125 | { 126 | dist_t diff = ax[i] - bx[i]; 127 | distance += diff * diff; 128 | } 129 | return sqrtf(distance); 130 | } 131 | #endif 132 | 133 | static dist_t cosine_dist_impl(coord_t const* ax, coord_t const* bx, size_t dim) 134 | { 135 | dist_t distance = 0.0; 136 | dist_t norma = 0.0; 137 | dist_t normb = 0.0; 138 | for (size_t i = 0; i < dim; i++) 139 | { 140 | distance += ax[i] * bx[i]; 141 | norma += ax[i] * ax[i]; 142 | normb += bx[i] * bx[i]; 143 | } 144 | return 1 - (distance / sqrt(norma * normb)); 145 | } 146 | 147 | static dist_t manhattan_dist_impl(coord_t const* ax, coord_t const* bx, size_t dim) 148 | { 149 | dist_t distance = 0.0; 150 | for (size_t i = 0; i < dim; i++) 151 | { 152 | distance += fabs(ax[i] - bx[i]); 153 | } 154 | return distance; 155 | } 156 | 157 | static dist_t (*dist_func_table[3])(coord_t const* ax, coord_t const* bx, size_t size); 158 | 159 | void hnsw_init_dist_func(void) 160 | { 161 | #ifdef __x86_64__ 162 | dist_func_table[DIST_L2] = __builtin_cpu_supports("avx2") 163 | ? l2_dist_impl_avx2 : l2_dist_impl_sse; 164 | #else 165 | dist_func_table[DIST_L2] = l2_dist_impl; 166 | #endif 167 | dist_func_table[DIST_COSINE] = cosine_dist_impl; 168 | dist_func_table[DIST_MANHATTAN] = manhattan_dist_impl; 169 | }; 170 | 171 | dist_t hnsw_dist_func(dist_func_t dist_func, coord_t const* ax, coord_t const* bx, size_t dim) 172 | { 173 | return dist_func_table[dist_func](ax, bx, dim); 174 | } 175 | -------------------------------------------------------------------------------- /embedding--0.3.5--0.3.6.sql: -------------------------------------------------------------------------------- 1 | -- Copyright 2023 Neon Inc. 2 | -- 3 | -- Licensed under the Apache License, Version 2.0 (the "License"); 4 | -- you may not use this file except in compliance with the License. 5 | -- You may obtain a copy of the License at 6 | -- 7 | -- http://www.apache.org/licenses/LICENSE-2.0 8 | -- 9 | -- Unless required by applicable law or agreed to in writing, software 10 | -- distributed under the License is distributed on an "AS IS" BASIS, 11 | -- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | -- See the License for the specific language governing permissions and 13 | -- limitations under the License. 14 | 15 | -- complain if script is sourced in psql, rather than via CREATE EXTENSION 16 | \echo Use "ALTER EXTENSION embedding UPDATE TO '0.3.6'" to load this file. \quit 17 | -------------------------------------------------------------------------------- /embedding--0.3.5.sql: -------------------------------------------------------------------------------- 1 | -- Copyright 2023 Neon Inc. 2 | -- 3 | -- Licensed under the Apache License, Version 2.0 (the "License"); 4 | -- you may not use this file except in compliance with the License. 5 | -- You may obtain a copy of the License at 6 | -- 7 | -- http://www.apache.org/licenses/LICENSE-2.0 8 | -- 9 | -- Unless required by applicable law or agreed to in writing, software 10 | -- distributed under the License is distributed on an "AS IS" BASIS, 11 | -- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | -- See the License for the specific language governing permissions and 13 | -- limitations under the License. 14 | 15 | -- complain if script is sourced in psql, rather than via CREATE EXTENSION 16 | \echo Use "CREATE EXTENSION embedding" to load this file. \quit 17 | 18 | -- functions 19 | 20 | CREATE FUNCTION l2_distance(real[], real[]) RETURNS real 21 | AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; 22 | 23 | CREATE FUNCTION cosine_distance(real[], real[]) RETURNS real 24 | AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; 25 | 26 | CREATE FUNCTION manhattan_distance(real[], real[]) RETURNS real 27 | AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; 28 | 29 | -- operators 30 | 31 | CREATE OPERATOR <-> ( 32 | LEFTARG = real[], RIGHTARG = real[], PROCEDURE = l2_distance, 33 | COMMUTATOR = '<->' 34 | ); 35 | 36 | CREATE OPERATOR <=> ( 37 | LEFTARG = real[], RIGHTARG = real[], PROCEDURE = cosine_distance, 38 | COMMUTATOR = '<=>' 39 | ); 40 | 41 | CREATE OPERATOR <~> ( 42 | LEFTARG = real[], RIGHTARG = real[], PROCEDURE = manhattan_distance, 43 | COMMUTATOR = '<~>' 44 | ); 45 | 46 | -- access method 47 | 48 | CREATE FUNCTION hnsw_handler(internal) RETURNS index_am_handler 49 | AS 'MODULE_PATHNAME' LANGUAGE C; 50 | 51 | CREATE ACCESS METHOD hnsw TYPE INDEX HANDLER hnsw_handler; 52 | 53 | COMMENT ON ACCESS METHOD hnsw IS 'hnsw index access method'; 54 | 55 | -- opclasses 56 | 57 | CREATE OPERATOR CLASS ann_l2_ops 58 | DEFAULT FOR TYPE real[] USING hnsw AS 59 | OPERATOR 1 <-> (real[], real[]) FOR ORDER BY float_ops, 60 | FUNCTION 1 l2_distance(real[], real[]); 61 | 62 | CREATE OPERATOR CLASS ann_cos_ops 63 | FOR TYPE real[] USING hnsw AS 64 | OPERATOR 1 <=> (real[], real[]) FOR ORDER BY float_ops, 65 | FUNCTION 1 cosine_distance(real[], real[]); 66 | 67 | CREATE OPERATOR CLASS ann_manhattan_ops 68 | FOR TYPE real[] USING hnsw AS 69 | OPERATOR 1 <~> (real[], real[]) FOR ORDER BY float_ops, 70 | FUNCTION 1 manhattan_distance(real[], real[]); 71 | -------------------------------------------------------------------------------- /embedding--0.3.6.sql: -------------------------------------------------------------------------------- 1 | -- Copyright 2023 Neon Inc. 2 | -- 3 | -- Licensed under the Apache License, Version 2.0 (the "License"); 4 | -- you may not use this file except in compliance with the License. 5 | -- You may obtain a copy of the License at 6 | -- 7 | -- http://www.apache.org/licenses/LICENSE-2.0 8 | -- 9 | -- Unless required by applicable law or agreed to in writing, software 10 | -- distributed under the License is distributed on an "AS IS" BASIS, 11 | -- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | -- See the License for the specific language governing permissions and 13 | -- limitations under the License. 14 | 15 | -- complain if script is sourced in psql, rather than via CREATE EXTENSION 16 | \echo Use "CREATE EXTENSION embedding" to load this file. \quit 17 | 18 | -- functions 19 | 20 | CREATE FUNCTION l2_distance(real[], real[]) RETURNS real 21 | AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; 22 | 23 | CREATE FUNCTION cosine_distance(real[], real[]) RETURNS real 24 | AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; 25 | 26 | CREATE FUNCTION manhattan_distance(real[], real[]) RETURNS real 27 | AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; 28 | 29 | -- operators 30 | 31 | CREATE OPERATOR <-> ( 32 | LEFTARG = real[], RIGHTARG = real[], PROCEDURE = l2_distance, 33 | COMMUTATOR = '<->' 34 | ); 35 | 36 | CREATE OPERATOR <=> ( 37 | LEFTARG = real[], RIGHTARG = real[], PROCEDURE = cosine_distance, 38 | COMMUTATOR = '<=>' 39 | ); 40 | 41 | CREATE OPERATOR <~> ( 42 | LEFTARG = real[], RIGHTARG = real[], PROCEDURE = manhattan_distance, 43 | COMMUTATOR = '<~>' 44 | ); 45 | 46 | -- access method 47 | 48 | CREATE FUNCTION hnsw_handler(internal) RETURNS index_am_handler 49 | AS 'MODULE_PATHNAME' LANGUAGE C; 50 | 51 | CREATE ACCESS METHOD hnsw TYPE INDEX HANDLER hnsw_handler; 52 | 53 | COMMENT ON ACCESS METHOD hnsw IS 'hnsw index access method'; 54 | 55 | -- opclasses 56 | 57 | CREATE OPERATOR CLASS ann_l2_ops 58 | DEFAULT FOR TYPE real[] USING hnsw AS 59 | OPERATOR 1 <-> (real[], real[]) FOR ORDER BY float_ops, 60 | FUNCTION 1 l2_distance(real[], real[]); 61 | 62 | CREATE OPERATOR CLASS ann_cos_ops 63 | FOR TYPE real[] USING hnsw AS 64 | OPERATOR 1 <=> (real[], real[]) FOR ORDER BY float_ops, 65 | FUNCTION 1 cosine_distance(real[], real[]); 66 | 67 | CREATE OPERATOR CLASS ann_manhattan_ops 68 | FOR TYPE real[] USING hnsw AS 69 | OPERATOR 1 <~> (real[], real[]) FOR ORDER BY float_ops, 70 | FUNCTION 1 manhattan_distance(real[], real[]); 71 | -------------------------------------------------------------------------------- /embedding.c: -------------------------------------------------------------------------------- 1 | // Copyright 2023 Neon Inc. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #include "postgres.h" 16 | 17 | #include "access/amapi.h" 18 | #include "access/generic_xlog.h" 19 | #include "access/relation.h" 20 | #include "access/reloptions.h" 21 | #include "access/tableam.h" 22 | #include "catalog/index.h" 23 | #include "commands/vacuum.h" 24 | #include "nodes/execnodes.h" 25 | #include "nodes/pathnodes.h" 26 | #include "storage/bufmgr.h" 27 | #include "storage/smgr.h" 28 | #include "utils/guc.h" 29 | #include "utils/selfuncs.h" 30 | #include "utils/spccache.h" 31 | 32 | #include 33 | #include 34 | 35 | #include "embedding.h" 36 | 37 | PG_MODULE_MAGIC; 38 | 39 | #define HNSW_DISTANCE_PROC 1 40 | #define HNSW_STACK_SIZE 4 41 | #define FIRST_PAGE 0 42 | 43 | /* Label flags */ 44 | #define DELETED_FLAG 1 45 | 46 | PGDLLEXPORT PG_FUNCTION_INFO_V1(l2_distance); 47 | PGDLLEXPORT PG_FUNCTION_INFO_V1(cosine_distance); 48 | PGDLLEXPORT PG_FUNCTION_INFO_V1(manhattan_distance); 49 | 50 | typedef union { 51 | label_t label; 52 | struct { 53 | ItemPointerData tid; 54 | uint16 flags; 55 | } pg; 56 | } HnswLabel; 57 | 58 | /* 59 | * Postgres specific part of HNSW index. 60 | * We are not poersisting this data, but reconstruct metadata from relation options. 61 | * There is not protectionf from altering index option for existed index, 62 | * butinfoirmation stored in opaque part of HNSW page allows to check if critical 63 | * metadata fields are changed (dimensiopns and maxM). 64 | */ 65 | typedef struct { 66 | HnswMetadata meta; 67 | Relation rel; 68 | bool unlogged; /* Do not need to wallog changes: either relation is unlogged, either index construction */ 69 | uint64_t n_inserted; /* Calculated since start of operation */ 70 | GenericXLogState* xlog_state; /* XLog state for wal logging updated pages */ 71 | Buffer writebuf; /* Currently written page */ 72 | Buffer lockbuf; /* First page is used to provide MURSIW access to HNSW index */ 73 | size_t n_buffers; /* Number of simultaneously accessed buffers */ 74 | Buffer buffers[HNSW_STACK_SIZE]; 75 | } HnswIndex; 76 | 77 | /* 78 | * This information in each HNSW page allows to detectincorrect metadata modification (ALTER INDEX) 79 | * which affects index format 80 | */ 81 | typedef struct 82 | { 83 | uint16_t dims; 84 | uint16_t maxM; 85 | } HnswPageOpaque; 86 | 87 | /* 88 | * Options associated with HNSW index, only "dims" is mandatory 89 | */ 90 | typedef struct { 91 | int32 vl_len_; /* varlena header (do not touch directly!) */ 92 | int dims; 93 | int efConstruction; 94 | int efSearch; 95 | int M; 96 | } HnswOptions; 97 | 98 | static relopt_kind hnsw_relopt_kind; 99 | 100 | typedef struct { 101 | HnswIndex* hnsw; 102 | size_t curr; 103 | size_t n_results; 104 | bool no_more_results; 105 | ArrayType* key; 106 | ItemPointer results; 107 | } HnswScanOpaqueData; 108 | 109 | typedef HnswScanOpaqueData* HnswScanOpaque; 110 | 111 | #define DEFAULT_EF_CONSTRUCT 16 112 | #define DEFAULT_EF_SEARCH 64 113 | #define DEFAULT_M 100 114 | 115 | static bool hnsw_add_point(HnswIndex* hnsw, coord_t const* coord, label_t label); 116 | 117 | PGDLLEXPORT void _PG_init(void); 118 | 119 | /* 120 | * Initialize index options and variables 121 | */ 122 | void 123 | _PG_init(void) 124 | { 125 | hnsw_relopt_kind = add_reloption_kind(); 126 | add_int_reloption(hnsw_relopt_kind, "dims", "Number of dimensions", 127 | 0, 0, INT_MAX 128 | #if PG_VERSION_NUM >= 130000 129 | , AccessExclusiveLock 130 | #endif 131 | ); 132 | add_int_reloption(hnsw_relopt_kind, "m", "Number of neighbors of each vertex", 133 | DEFAULT_M, 0, INT_MAX 134 | #if PG_VERSION_NUM >= 130000 135 | , AccessExclusiveLock 136 | #endif 137 | ); 138 | add_int_reloption(hnsw_relopt_kind, "efconstruction", "Number of inspected neighbors during index construction", 139 | DEFAULT_EF_CONSTRUCT, 1, INT_MAX 140 | #if PG_VERSION_NUM >= 130000 141 | , AccessExclusiveLock 142 | #endif 143 | ); 144 | add_int_reloption(hnsw_relopt_kind, "efsearch", "Number of inspected neighbors during index search", 145 | DEFAULT_EF_SEARCH, 1, INT_MAX 146 | #if PG_VERSION_NUM >= 130000 147 | , AccessExclusiveLock 148 | #endif 149 | ); 150 | hnsw_init_dist_func(); 151 | } 152 | 153 | static void 154 | hnsw_build_callback(Relation index, 155 | #if PG_VERSION_NUM >= 130000 156 | ItemPointer tid, 157 | #else 158 | HeapTuple hup, 159 | #endif 160 | Datum *values, bool *isnull, bool tupleIsAlive, void *state) 161 | { 162 | HnswIndex* hnsw = (HnswIndex*) state; 163 | ArrayType* array; 164 | int n_items; 165 | HnswLabel u; 166 | 167 | #if PG_VERSION_NUM < 130000 168 | ItemPointer tid = &hup->t_self; 169 | #endif 170 | 171 | /* Skip nulls */ 172 | if (isnull[0]) 173 | return; 174 | 175 | array = DatumGetArrayTypePCopy(values[0]); 176 | n_items = ArrayGetNItems(ARR_NDIM(array), ARR_DIMS(array)); 177 | if (n_items != hnsw->meta.dim) 178 | { 179 | elog(ERROR, "Wrong number of dimensions: %d instead of %d expected", 180 | n_items, (int)hnsw->meta.dim); 181 | } 182 | 183 | u.pg.tid = *tid; 184 | u.pg.flags = 0; 185 | 186 | if (!hnsw_add_point(hnsw, (coord_t*)ARR_DATA_PTR(array), u.label)) 187 | elog(ERROR, "HNSW index insert failed"); 188 | pfree(array); 189 | } 190 | 191 | static dist_func_t 192 | hnsw_resolve_dist_func(Relation index) 193 | { 194 | FmgrInfo* proc_info = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); 195 | if (proc_info->fn_addr == l2_distance) 196 | return DIST_L2; 197 | else if (proc_info->fn_addr == cosine_distance) 198 | return DIST_COSINE; 199 | else if (proc_info->fn_addr == manhattan_distance) 200 | return DIST_MANHATTAN; 201 | else 202 | elog(ERROR, "Function is not supported by HNSW inodex"); 203 | } 204 | 205 | static void 206 | hnsw_populate(HnswIndex* hnsw, Relation indexRel, Relation heapRel) 207 | { 208 | IndexInfo* indexInfo = BuildIndexInfo(indexRel); 209 | Assert(indexInfo->ii_NumIndexAttrs == 1); 210 | table_index_build_scan(heapRel, indexRel, indexInfo, 211 | true, true, hnsw_build_callback, (void *)hnsw, NULL); 212 | } 213 | 214 | static HnswIndex* 215 | hnsw_get_index(Relation indexRel) 216 | { 217 | HnswIndex* hnsw = (HnswIndex*)palloc(sizeof(HnswIndex)); 218 | HnswOptions *opts = (HnswOptions *) indexRel->rd_options; 219 | if (opts == NULL || opts->dims == 0) { 220 | elog(ERROR, "HNSW index requires 'dims' to be specified"); 221 | } 222 | hnsw->meta.dim = opts->dims; 223 | hnsw->meta.M = opts->M; 224 | hnsw->meta.maxM = hnsw->meta.M * 2; 225 | hnsw->meta.data_size = hnsw->meta.dim * sizeof(coord_t); 226 | hnsw->meta.offset_data = (hnsw->meta.maxM + 1) * sizeof(idx_t); 227 | hnsw->meta.offset_label = hnsw->meta.offset_data + hnsw->meta.data_size; 228 | hnsw->meta.size_data_per_element = hnsw->meta.offset_label + sizeof(label_t); 229 | hnsw->meta.elems_per_page = (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - sizeof(HnswPageOpaque)) / (hnsw->meta.size_data_per_element + sizeof(ItemIdData)); 230 | if (hnsw->meta.elems_per_page == 0) 231 | elog(ERROR, "Element doesn't fit in Postgres page"); 232 | hnsw->meta.efConstruction = opts->efConstruction; 233 | hnsw->meta.efSearch = opts->efSearch; 234 | hnsw->meta.dist_func = hnsw_resolve_dist_func(indexRel); 235 | hnsw->meta.enterpoint_node = 0; 236 | hnsw->rel = indexRel; 237 | hnsw->n_buffers = 0; 238 | hnsw->xlog_state = NULL; 239 | hnsw->n_inserted = 0; 240 | hnsw->unlogged = RelationNeedsWAL(indexRel); 241 | hnsw->lockbuf = InvalidBuffer; 242 | hnsw->writebuf = InvalidBuffer; 243 | return hnsw; 244 | } 245 | 246 | /* 247 | * Start or restart an index scan 248 | */ 249 | static IndexScanDesc 250 | hnsw_beginscan(Relation index, int nkeys, int norderbys) 251 | { 252 | IndexScanDesc scan = RelationGetIndexScan(index, nkeys, norderbys); 253 | HnswScanOpaque so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData)); 254 | so->hnsw = hnsw_get_index(index); 255 | so->curr = 0; 256 | so->n_results = 0; 257 | so->results = NULL; 258 | so->no_more_results = true; 259 | so->key = NULL; 260 | scan->opaque = so; 261 | return scan; 262 | } 263 | 264 | /* 265 | * Start or restart an index scan 266 | */ 267 | static void 268 | hnsw_rescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys) 269 | { 270 | HnswScanOpaque so = (HnswScanOpaque) scan->opaque; 271 | if (so->results) 272 | { 273 | pfree(so->results); 274 | so->results = NULL; 275 | } 276 | so->curr = 0; 277 | if (orderbys && scan->numberOfOrderBys > 0) 278 | memmove(scan->orderByData, orderbys, scan->numberOfOrderBys * sizeof(ScanKeyData)); 279 | } 280 | 281 | /* 282 | * Fetch the next tuple in the given scan 283 | */ 284 | static bool 285 | hnsw_gettuple(IndexScanDesc scan, ScanDirection dir) 286 | { 287 | HnswScanOpaque so = (HnswScanOpaque) scan->opaque; 288 | size_t n_results; 289 | label_t* results; 290 | 291 | /* 292 | * Index can be used to scan backward, but Postgres doesn't support 293 | * backward scan on operators 294 | */ 295 | Assert(ScanDirectionIsForward(dir)); 296 | 297 | if (so->curr == 0) 298 | { 299 | Datum value; 300 | int n_items; 301 | 302 | /* Safety check */ 303 | if (scan->orderByData == NULL) 304 | elog(ERROR, "cannot scan HNSW index without order"); 305 | 306 | /* No items will match if null */ 307 | if (scan->orderByData->sk_flags & SK_ISNULL) 308 | return false; 309 | 310 | value = scan->orderByData->sk_argument; 311 | so->key = DatumGetArrayTypePCopy(value); 312 | n_items = ArrayGetNItems(ARR_NDIM(so->key), ARR_DIMS(so->key)); 313 | if (n_items != so->hnsw->meta.dim) 314 | elog(ERROR, "Wrong number of dimensions: %d instead of %d expected", 315 | n_items, (int)so->hnsw->meta.dim); 316 | 317 | if (!hnsw_search(&so->hnsw->meta, (coord_t*)ARR_DATA_PTR(so->key), &n_results, &results)) 318 | elog(ERROR, "HNSW index search failed"); 319 | 320 | so->results = (ItemPointer)palloc(n_results*sizeof(ItemPointerData)); 321 | so->n_results = n_results; 322 | so->no_more_results = n_results < so->hnsw->meta.efSearch; 323 | for (size_t i = 0; i < n_results; i++) 324 | { 325 | memcpy(&so->results[i], &results[i], sizeof(so->results[i])); 326 | } 327 | free(results); 328 | } 329 | if (so->curr >= so->n_results) 330 | { 331 | if (so->no_more_results) 332 | return false; 333 | 334 | so->hnsw->meta.efSearch *= 2; 335 | if (!hnsw_search(&so->hnsw->meta, (coord_t*)ARR_DATA_PTR(so->key), &n_results, &results)) 336 | elog(ERROR, "HNSW index search failed"); 337 | 338 | if (n_results <= so->n_results) 339 | { 340 | /* No new results found */ 341 | return false; 342 | } 343 | so->no_more_results = n_results < so->hnsw->meta.efSearch; 344 | 345 | /* ANN search with larger K (efSearch) can find better results than with smaller K. 346 | * We have two choices: 347 | * 1. Ignore them to preserve monotony of results. 348 | * 2. Include them to include more relevant results in selection and increase recall 349 | * To ignore them we need hnsw_search to also return distance. 350 | * Without it the only choice is 2) 351 | */ 352 | so->results = (ItemPointer)repalloc(so->results, (n_results + so->n_results)*sizeof(ItemPointerData)); 353 | 354 | /* Sort for binary search */ 355 | pg_qsort(so->results, so->n_results, sizeof(ItemPointerData), (int (*)(const void *, const void *))ItemPointerCompare); 356 | 357 | /* Exclude already returned records */ 358 | for (size_t i = 0; i < n_results; i++) 359 | { 360 | if (!bsearch(&results[i], so->results, so->n_results, sizeof(ItemPointerData), (int (*)(const void *, const void *))ItemPointerCompare)) 361 | { 362 | memcpy(&so->results[so->n_results++], &results[i], sizeof(ItemPointerData)); 363 | } 364 | } 365 | Assert(so->curr < so->n_results); 366 | } 367 | scan->xs_heaptid = so->results[so->curr++]; 368 | scan->xs_recheckorderby = false; 369 | return true; 370 | } 371 | 372 | /* 373 | * End a scan and release resources 374 | */ 375 | static void 376 | hnsw_endscan(IndexScanDesc scan) 377 | { 378 | HnswScanOpaque so = (HnswScanOpaque) scan->opaque; 379 | if (so->key) 380 | pfree(so->key); 381 | if (so->results) 382 | pfree(so->results); 383 | if (so->hnsw) 384 | pfree(so->hnsw); 385 | pfree(so); 386 | scan->opaque = NULL; 387 | } 388 | 389 | 390 | /* 391 | * Estimate the cost of an index scan 392 | */ 393 | static void 394 | hnsw_costestimate(PlannerInfo *root, IndexPath *path, double loop_count, 395 | Cost *indexStartupCost, Cost *indexTotalCost, 396 | Selectivity *indexSelectivity, double *indexCorrelation 397 | ,double *indexPages 398 | ) 399 | { 400 | GenericCosts costs; 401 | 402 | /* Never use index without order */ 403 | if (path->indexorderbys == NULL) 404 | { 405 | *indexStartupCost = DBL_MAX; 406 | *indexTotalCost = DBL_MAX; 407 | *indexSelectivity = 0; 408 | *indexCorrelation = 0; 409 | *indexPages = 0; 410 | return; 411 | } 412 | else 413 | { 414 | IndexOptInfo *index = path->indexinfo; 415 | Relation rel = index_open(index->indexoid, NoLock); 416 | HnswIndex* hnsw = hnsw_get_index(rel); 417 | double spc_random_page_cost; 418 | 419 | get_tablespace_page_costs(index->reltablespace, 420 | &spc_random_page_cost, 421 | NULL); 422 | 423 | MemSet(&costs, 0, sizeof(costs)); 424 | 425 | genericcostestimate(root, path, loop_count, &costs); 426 | 427 | /* Number of pages inspected by search is limited by efSearch parameter */ 428 | *indexStartupCost = *indexTotalCost = hnsw->meta.efSearch * spc_random_page_cost; 429 | *indexSelectivity = index->rel->rows ? hnsw->meta.efSearch / index->rel->rows : costs.indexSelectivity; 430 | *indexCorrelation = costs.indexCorrelation; 431 | *indexPages = hnsw->meta.efSearch; 432 | 433 | pfree(hnsw); 434 | index_close(rel, NoLock); 435 | } 436 | } 437 | 438 | /* 439 | * Parse and validate the reloptions 440 | */ 441 | static bytea * 442 | hnsw_options(Datum reloptions, bool validate) 443 | { 444 | static const relopt_parse_elt tab[] = { 445 | {"dims", RELOPT_TYPE_INT, offsetof(HnswOptions, dims)}, 446 | {"efconstruction", RELOPT_TYPE_INT, offsetof(HnswOptions, efConstruction)}, 447 | {"efsearch", RELOPT_TYPE_INT, offsetof(HnswOptions, efSearch)}, 448 | {"m", RELOPT_TYPE_INT, offsetof(HnswOptions, M)} 449 | }; 450 | 451 | #if PG_VERSION_NUM >= 130000 452 | return (bytea *) build_reloptions(reloptions, validate, 453 | hnsw_relopt_kind, 454 | sizeof(HnswOptions), 455 | tab, lengthof(tab)); 456 | #else 457 | relopt_value *options; 458 | HnswOptions *rdopts; 459 | int numoptions; 460 | 461 | options = parseRelOptions(reloptions, validate, hnsw_relopt_kind, &numoptions); 462 | 463 | rdopts = allocateReloptStruct(sizeof(HnswOptions), options, numoptions); 464 | 465 | fillRelOptions((void *) rdopts, sizeof(HnswOptions), options, numoptions, validate, tab, lengthof(tab)); 466 | 467 | return (bytea *) rdopts; 468 | #endif 469 | } 470 | 471 | /* 472 | * Validate catalog entries for the specified operator class 473 | */ 474 | static bool 475 | hnsw_validate(Oid opclassoid) 476 | { 477 | return true; 478 | } 479 | 480 | 481 | /* We need to initialize firtst page to avoid race condition on insert */ 482 | static void hnsw_init_first_page(HnswIndex* hnsw, ForkNumber forknum) 483 | { 484 | Buffer buf; 485 | HnswPageOpaque* opq; 486 | Page page; 487 | 488 | buf = ReadBufferExtended(hnsw->rel, forknum, P_NEW, RBM_NORMAL, NULL); 489 | Assert(BufferGetBlockNumber(buf) == FIRST_PAGE); 490 | LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); 491 | page = BufferGetPage(buf); 492 | PageInit(page, BufferGetPageSize(buf), sizeof(HnswPageOpaque)); 493 | opq = (HnswPageOpaque*)PageGetSpecialPointer(page); 494 | opq->dims = (uint16_t)hnsw->meta.dim; 495 | opq->maxM = (uint16_t)hnsw->meta.maxM; 496 | MarkBufferDirty(buf); 497 | UnlockReleaseBuffer(buf); 498 | } 499 | 500 | /* 501 | * Build the index for a logged table 502 | */ 503 | static IndexBuildResult * 504 | hnsw_build(Relation heap, Relation index, IndexInfo *indexInfo) 505 | { 506 | HnswIndex* hnsw = hnsw_get_index(index); 507 | IndexBuildResult* result = (IndexBuildResult *) palloc(sizeof(IndexBuildResult)); 508 | 509 | hnsw->unlogged = true; 510 | #ifdef NEON_SMGR 511 | smgr_start_unlogged_build(RelationGetSmgr(index)); 512 | #endif 513 | 514 | hnsw_init_first_page(hnsw, MAIN_FORKNUM); 515 | 516 | hnsw_populate(hnsw, index, heap); 517 | 518 | #ifdef NEON_SMGR 519 | smgr_finish_unlogged_build_phase_1(RelationGetSmgr(index)); 520 | #endif 521 | 522 | /* 523 | * We didn't write WAL records as we built the index, so if 524 | * WAL-logging is required, write all pages to the WAL now. 525 | */ 526 | if (RelationNeedsWAL(index)) 527 | { 528 | log_newpage_range(index, MAIN_FORKNUM, 529 | 0, RelationGetNumberOfBlocks(index), 530 | true); 531 | #ifdef NEON_SMGR 532 | { 533 | #if PG_VERSION_NUM >= 160000 534 | RelFileLocator locator = index->rd_locator; 535 | #else 536 | RelFileNode locator = index->rd_node; 537 | #endif 538 | SetLastWrittenLSNForBlockRange(XactLastRecEnd, locator, 539 | MAIN_FORKNUM, 0, RelationGetNumberOfBlocks(index)); 540 | SetLastWrittenLSNForRelation(XactLastRecEnd, locator, MAIN_FORKNUM); 541 | } 542 | #endif 543 | } 544 | #ifdef NEON_SMGR 545 | smgr_end_unlogged_build(RelationGetSmgr(index)); 546 | #endif 547 | 548 | result->heap_tuples = result->index_tuples = hnsw->n_inserted; 549 | pfree(hnsw); 550 | return result; 551 | } 552 | 553 | /* 554 | * Insert a tuple into the index 555 | */ 556 | static bool 557 | hnsw_insert(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, 558 | Relation heap, IndexUniqueCheck checkUnique, 559 | #if PG_VERSION_NUM >= 140000 560 | bool indexUnchanged, 561 | #endif 562 | IndexInfo *indexInfo) 563 | { 564 | ArrayType* array; 565 | int n_items; 566 | HnswLabel u; 567 | HnswIndex* hnsw; 568 | bool success; 569 | 570 | /* Skip nulls */ 571 | if (isnull[0]) 572 | return false; 573 | 574 | hnsw = hnsw_get_index(index); 575 | 576 | /* Detoast value */ 577 | array = DatumGetArrayTypePCopy(values[0]); 578 | n_items = ArrayGetNItems(ARR_NDIM(array), ARR_DIMS(array)); 579 | if (n_items != hnsw->meta.dim) 580 | { 581 | elog(ERROR, "Wrong number of dimensions: %d instead of %d expected", 582 | n_items, (int)hnsw->meta.dim); 583 | } 584 | 585 | u.pg.tid = *heap_tid; 586 | u.pg.flags = 0; 587 | 588 | success = hnsw_add_point(hnsw, (coord_t*)ARR_DATA_PTR(array), u.label); 589 | pfree(array); 590 | pfree(hnsw); 591 | return success; 592 | } 593 | 594 | static void hnsw_check_meta(HnswMetadata* meta, Page page) 595 | { 596 | HnswPageOpaque* opq = (HnswPageOpaque*)PageGetSpecialPointer(page); 597 | if (opq->dims != (uint16_t)meta->dim || 598 | opq->maxM != (uint16_t)meta->maxM) 599 | { 600 | elog(ERROR, "Inconsistency with HNSW index metadata: only ef_construction and ef_search options of HNSW index may be altered"); 601 | } 602 | } 603 | 604 | 605 | 606 | static bool hnsw_add_point(HnswIndex* hnsw, coord_t const* coord, label_t label) 607 | { 608 | BlockNumber rel_size; 609 | GenericXLogState *state = NULL; 610 | OffsetNumber ins_offs = 0; 611 | bool extend = false; 612 | idx_t cur_c; 613 | Buffer buf; 614 | Page page; 615 | HnswPageOpaque* opq; 616 | bool result; 617 | char item[BLCKSZ]; 618 | 619 | memset(item, 0, hnsw->meta.offset_data); 620 | memcpy(item + hnsw->meta.offset_data, coord, hnsw->meta.offset_label - hnsw->meta.offset_data); 621 | memcpy(item + hnsw->meta.offset_label, &label, sizeof(label_t)); 622 | 623 | 624 | /* First obtain exclusive lock oni first page: itis needed to sycnhronize access to the index. 625 | * We done not support concurrent inserted because HNSW insert algorithm can acess pages in arbitrary ortder and sop cause deadlock 626 | */ 627 | Assert(hnsw->lockbuf == InvalidBuffer); 628 | hnsw->lockbuf = ReadBuffer(hnsw->rel, FIRST_PAGE); 629 | LockBuffer(hnsw->lockbuf, BUFFER_LOCK_EXCLUSIVE); 630 | /* Obtain size under lock */ 631 | rel_size = RelationGetNumberOfBlocks(hnsw->rel); 632 | 633 | while (true) 634 | { 635 | if (extend) 636 | { 637 | buf = ReadBuffer(hnsw->rel, P_NEW); 638 | LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); 639 | } 640 | else 641 | { 642 | if (rel_size-1 != FIRST_PAGE) 643 | { 644 | buf = ReadBuffer(hnsw->rel, rel_size - 1); 645 | LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); 646 | } 647 | else 648 | buf = hnsw->lockbuf; 649 | } 650 | 651 | if (!hnsw->unlogged) 652 | state = GenericXLogStart(hnsw->rel); 653 | 654 | page = hnsw->unlogged ? BufferGetPage(buf) : GenericXLogRegisterBuffer(state, buf, extend ? GENERIC_XLOG_FULL_IMAGE : 0); 655 | if (extend) 656 | { 657 | Assert(BufferGetBlockNumber(buf) == rel_size); 658 | PageInit(page, BufferGetPageSize(buf), sizeof(HnswPageOpaque)); 659 | opq = (HnswPageOpaque*)PageGetSpecialPointer(page); 660 | opq->dims = (uint16_t)hnsw->meta.dim; 661 | opq->maxM = (uint16_t)hnsw->meta.maxM; 662 | rel_size += 1; 663 | } 664 | else 665 | { 666 | if (hnsw->lockbuf == buf) 667 | hnsw_check_meta(&hnsw->meta, page); 668 | } 669 | 670 | ins_offs = PageAddItem(page, (Item)item, hnsw->meta.size_data_per_element, InvalidOffsetNumber, false, false); 671 | if (ins_offs == InvalidOffsetNumber) 672 | { 673 | if (extend) 674 | elog(ERROR, "Failed to append item to the page"); 675 | if (state) 676 | GenericXLogAbort(state); 677 | if (buf != hnsw->lockbuf) 678 | UnlockReleaseBuffer(buf); 679 | extend = true; 680 | } 681 | else 682 | break; 683 | } 684 | MarkBufferDirty(buf); 685 | if (state) 686 | GenericXLogFinish(state); 687 | 688 | if (buf != hnsw->lockbuf) 689 | UnlockReleaseBuffer(buf); 690 | 691 | hnsw->n_inserted += 1; 692 | 693 | cur_c = (rel_size-1)*hnsw->meta.elems_per_page + ins_offs - FirstOffsetNumber; 694 | 695 | result = hnsw_bind_point(&hnsw->meta, coord, cur_c); 696 | 697 | UnlockReleaseBuffer(hnsw->lockbuf); 698 | hnsw->lockbuf = InvalidBuffer; 699 | 700 | return result; 701 | } 702 | 703 | 704 | bool hnsw_begin_read(HnswMetadata* meta, idx_t idx, idx_t** indexes, coord_t** coords, label_t* label) 705 | { 706 | HnswIndex* hnsw = (HnswIndex*)meta; 707 | BlockNumber blkno = idx/meta->elems_per_page; 708 | Page page; 709 | Item item; 710 | ItemId item_id; 711 | OffsetNumber offset; 712 | Buffer buf; 713 | 714 | if (hnsw->n_buffers >= HNSW_STACK_SIZE) 715 | elog(ERROR, "HNSW stack overflow"); 716 | 717 | /* First page is already locked for exclusive update of index */ 718 | if (blkno == FIRST_PAGE && hnsw->lockbuf != InvalidBuffer) 719 | { 720 | buf = hnsw->lockbuf; 721 | } 722 | else if (hnsw->writebuf != InvalidBuffer && BufferGetBlockNumber(hnsw->writebuf) == blkno) 723 | { 724 | buf = hnsw->writebuf; 725 | } 726 | else 727 | { 728 | buf = ReadBuffer(hnsw->rel, blkno); 729 | LockBuffer(buf, BUFFER_LOCK_SHARE); 730 | } 731 | page = BufferGetPage(buf); 732 | 733 | if (blkno == FIRST_PAGE) 734 | hnsw_check_meta(meta, page); 735 | 736 | offset = FirstOffsetNumber + idx % meta->elems_per_page; 737 | if (offset > PageGetMaxOffsetNumber(page)) 738 | { 739 | if (buf != hnsw->lockbuf && buf != hnsw->writebuf) 740 | UnlockReleaseBuffer(buf); 741 | return false; 742 | } 743 | item_id = PageGetItemId(page, offset); 744 | item = PageGetItem(page, item_id); 745 | 746 | hnsw->buffers[hnsw->n_buffers++] = buf; 747 | 748 | if (indexes) 749 | *indexes = (idx_t*)item; 750 | 751 | if (coords) 752 | *coords = (coord_t*)((char*)item + meta->offset_data); 753 | 754 | if (label) 755 | memcpy(label, (char*)item + meta->offset_label, sizeof(*label)); 756 | return true; 757 | } 758 | 759 | void hnsw_end_read(HnswMetadata* meta) 760 | { 761 | HnswIndex* hnsw = (HnswIndex*)meta; 762 | if (hnsw->n_buffers == 0) 763 | elog(ERROR, "HNSW stack is empty"); 764 | hnsw->n_buffers -= 1; 765 | if (hnsw->buffers[hnsw->n_buffers] != hnsw->lockbuf && hnsw->buffers[hnsw->n_buffers] != hnsw->writebuf) 766 | UnlockReleaseBuffer(hnsw->buffers[hnsw->n_buffers]); 767 | } 768 | 769 | void hnsw_begin_write(HnswMetadata* meta, idx_t idx, idx_t** indexes, coord_t** coords, label_t* label) 770 | { 771 | HnswIndex* hnsw = (HnswIndex*)meta; 772 | BlockNumber blkno = idx/meta->elems_per_page; 773 | Page page; 774 | ItemId item_id; 775 | Item item; 776 | Buffer buf; 777 | 778 | Assert(hnsw->lockbuf != InvalidBuffer); /* index should be exclsuively locked */ 779 | 780 | if (hnsw->xlog_state || hnsw->writebuf != InvalidBuffer) 781 | elog(ERROR, "More than two concurrent write operations"); 782 | 783 | if (hnsw->n_buffers >= HNSW_STACK_SIZE) 784 | elog(ERROR, "HNSW stack overflow"); 785 | 786 | /* First page is already locked for exclusive update of index */ 787 | if (blkno == FIRST_PAGE) 788 | { 789 | buf = hnsw->lockbuf; 790 | } 791 | else 792 | { 793 | buf = ReadBuffer(hnsw->rel, blkno); 794 | hnsw->writebuf = buf; 795 | LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); 796 | } 797 | if (hnsw->unlogged) 798 | { 799 | hnsw->xlog_state = NULL; 800 | page = BufferGetPage(buf); 801 | } 802 | else 803 | { 804 | hnsw->xlog_state = GenericXLogStart(hnsw->rel); 805 | page = GenericXLogRegisterBuffer(hnsw->xlog_state, buf, 0); 806 | } 807 | 808 | item_id = PageGetItemId(page, FirstOffsetNumber + idx % meta->elems_per_page); 809 | item = PageGetItem(page, item_id); 810 | 811 | hnsw->buffers[hnsw->n_buffers++] = buf; 812 | 813 | if (indexes) 814 | *indexes = (idx_t*)item; 815 | 816 | if (coords) 817 | *coords = (coord_t*)((char*)item + meta->offset_data); 818 | 819 | if (label) 820 | memcpy(label, (char*)item + meta->offset_label, sizeof(*label)); 821 | } 822 | 823 | void hnsw_end_write(HnswMetadata* meta) 824 | { 825 | HnswIndex* hnsw = (HnswIndex*)meta; 826 | 827 | if (hnsw->n_buffers == 0) 828 | elog(ERROR, "HNSW stack is empty"); 829 | 830 | hnsw->n_buffers -= 1; 831 | MarkBufferDirty(hnsw->buffers[hnsw->n_buffers]); 832 | if (hnsw->xlog_state) 833 | { 834 | GenericXLogFinish(hnsw->xlog_state); 835 | hnsw->xlog_state = NULL; 836 | } 837 | if (hnsw->buffers[hnsw->n_buffers] != hnsw->lockbuf) 838 | { 839 | Assert(hnsw->buffers[hnsw->n_buffers] == hnsw->writebuf); 840 | hnsw->writebuf = InvalidBuffer; 841 | UnlockReleaseBuffer(hnsw->buffers[hnsw->n_buffers]); 842 | } 843 | } 844 | 845 | void hnsw_prefetch(HnswMetadata* meta, idx_t idx) 846 | { 847 | HnswIndex* hnsw = (HnswIndex*)meta; 848 | BlockNumber blkno = idx/meta->elems_per_page; 849 | PrefetchBuffer(hnsw->rel, MAIN_FORKNUM, blkno); 850 | } 851 | 852 | 853 | /* 854 | * Build the index for an unlogged table 855 | */ 856 | static void 857 | hnsw_buildempty(Relation index) 858 | { 859 | HnswIndex* hnsw = hnsw_get_index(index); 860 | hnsw_init_first_page(hnsw, INIT_FORKNUM); 861 | pfree(hnsw); 862 | } 863 | 864 | /* 865 | * Clean up after a VACUUM operation 866 | */ 867 | static IndexBulkDeleteResult * 868 | hnsw_vacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats) 869 | { 870 | Relation rel = info->index; 871 | 872 | if (stats == NULL) 873 | return NULL; 874 | 875 | stats->num_pages = RelationGetNumberOfBlocks(rel); 876 | 877 | return stats; 878 | } 879 | 880 | /* 881 | * Bulk delete tuples from the index 882 | */ 883 | static IndexBulkDeleteResult * 884 | hnsw_bulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats, 885 | IndexBulkDeleteCallback callback, void *callback_state) 886 | { 887 | OffsetNumber maxoffno; 888 | OffsetNumber offno; 889 | Relation index = info->index; 890 | Buffer buf; 891 | Page page; 892 | int n_deleted; 893 | GenericXLogState *state; 894 | BlockNumber rel_size = RelationGetNumberOfBlocks(index); 895 | BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD); 896 | HnswIndex* hnsw = hnsw_get_index(index); 897 | 898 | if (stats == NULL) 899 | stats = (IndexBulkDeleteResult *) palloc0(sizeof(IndexBulkDeleteResult)); 900 | 901 | for (BlockNumber blkno = FIRST_PAGE; blkno < rel_size; blkno++) 902 | { 903 | buf = ReadBufferExtended(index, MAIN_FORKNUM, blkno, RBM_NORMAL, bas); 904 | 905 | /* 906 | * ambulkdelete cannot delete entries from pages that are 907 | * pinned by other backends 908 | * 909 | * https://www.postgresql.org/docs/current/index-locking.html 910 | */ 911 | LockBufferForCleanup(buf); 912 | state = GenericXLogStart(index); 913 | page = GenericXLogRegisterBuffer(state, buf, 0); 914 | 915 | maxoffno = PageGetMaxOffsetNumber(page); 916 | n_deleted = 0; 917 | 918 | for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) 919 | { 920 | HnswLabel* label = (HnswLabel*)((char*)PageGetItem(page, PageGetItemId(page, offno)) + hnsw->meta.offset_label); 921 | if (!(label->pg.flags & DELETED_FLAG)) 922 | { 923 | if (callback(&label->pg.tid, callback_state)) 924 | { 925 | label->pg.flags |= DELETED_FLAG; 926 | stats->tuples_removed++; 927 | n_deleted += 1; 928 | } 929 | else 930 | stats->num_index_tuples++; 931 | } 932 | } 933 | if (n_deleted > 0) 934 | { 935 | MarkBufferDirty(buf); 936 | GenericXLogFinish(state); 937 | } 938 | else 939 | GenericXLogAbort(state); 940 | 941 | UnlockReleaseBuffer(buf); 942 | } 943 | pfree(hnsw); 944 | 945 | return stats; 946 | } 947 | 948 | bool hnsw_is_deleted(label_t label) 949 | { 950 | HnswLabel u; 951 | u.label = label; 952 | return (u.pg.flags & DELETED_FLAG) != 0; 953 | } 954 | 955 | /* 956 | * Define index handler 957 | * 958 | * See https://www.postgresql.org/docs/current/index-api.html 959 | */ 960 | PGDLLEXPORT PG_FUNCTION_INFO_V1(hnsw_handler); 961 | Datum 962 | hnsw_handler(PG_FUNCTION_ARGS) 963 | { 964 | IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); 965 | 966 | amroutine->amstrategies = 0; 967 | amroutine->amsupport = 1; 968 | #if PG_VERSION_NUM >= 130000 969 | amroutine->amoptsprocnum = 0; 970 | #endif 971 | amroutine->amcanorder = false; 972 | amroutine->amcanorderbyop = true; 973 | amroutine->amcanbackward = false; /* can change direction mid-scan */ 974 | amroutine->amcanunique = false; 975 | amroutine->amcanmulticol = false; 976 | amroutine->amoptionalkey = true; 977 | amroutine->amsearcharray = false; 978 | amroutine->amsearchnulls = false; 979 | amroutine->amstorage = false; 980 | amroutine->amclusterable = false; 981 | amroutine->ampredlocks = false; 982 | amroutine->amcanparallel = false; 983 | amroutine->amcaninclude = false; 984 | #if PG_VERSION_NUM >= 130000 985 | amroutine->amusemaintenanceworkmem = false; /* not used during VACUUM */ 986 | amroutine->amparallelvacuumoptions = VACUUM_OPTION_PARALLEL_BULKDEL; 987 | #endif 988 | amroutine->amkeytype = InvalidOid; 989 | 990 | /* Interface functions */ 991 | amroutine->ambuild = hnsw_build; 992 | amroutine->ambuildempty = hnsw_buildempty; 993 | amroutine->aminsert = hnsw_insert; 994 | amroutine->ambulkdelete = hnsw_bulkdelete; 995 | amroutine->amvacuumcleanup = hnsw_vacuumcleanup; 996 | amroutine->amcanreturn = NULL; /* tuple not included in heapsort */ 997 | amroutine->amcostestimate = hnsw_costestimate; 998 | amroutine->amoptions = hnsw_options; 999 | amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */ 1000 | amroutine->ambuildphasename = NULL; 1001 | amroutine->amvalidate = hnsw_validate; 1002 | #if PG_VERSION_NUM >= 140000 1003 | amroutine->amadjustmembers = NULL; 1004 | #endif 1005 | amroutine->ambeginscan = hnsw_beginscan; 1006 | amroutine->amrescan = hnsw_rescan; 1007 | amroutine->amgettuple = hnsw_gettuple; 1008 | amroutine->amgetbitmap = NULL; 1009 | amroutine->amendscan = hnsw_endscan; 1010 | amroutine->ammarkpos = NULL; 1011 | amroutine->amrestrpos = NULL; 1012 | 1013 | /* Interface functions to support parallel index scans */ 1014 | amroutine->amestimateparallelscan = NULL; 1015 | amroutine->aminitparallelscan = NULL; 1016 | amroutine->amparallelrescan = NULL; 1017 | 1018 | PG_RETURN_POINTER(amroutine); 1019 | } 1020 | 1021 | 1022 | static dist_t 1023 | calc_distance(dist_func_t dist, ArrayType *a, ArrayType *b) 1024 | { 1025 | int a_dim = ArrayGetNItems(ARR_NDIM(a), ARR_DIMS(a)); 1026 | int b_dim = ArrayGetNItems(ARR_NDIM(b), ARR_DIMS(b)); 1027 | coord_t *ax = (coord_t*)ARR_DATA_PTR(a); 1028 | coord_t *bx = (coord_t*)ARR_DATA_PTR(b); 1029 | 1030 | if (a_dim != b_dim) 1031 | { 1032 | ereport(ERROR, 1033 | (errcode(ERRCODE_DATA_EXCEPTION), 1034 | errmsg("different array dimensions %d and %d", a_dim, b_dim))); 1035 | } 1036 | 1037 | return hnsw_dist_func(dist, ax, bx, a_dim); 1038 | } 1039 | 1040 | Datum 1041 | l2_distance(PG_FUNCTION_ARGS) 1042 | { 1043 | ArrayType *a = PG_GETARG_ARRAYTYPE_P(0); 1044 | ArrayType *b = PG_GETARG_ARRAYTYPE_P(1); 1045 | PG_RETURN_FLOAT4(calc_distance(DIST_L2, a, b)); 1046 | } 1047 | 1048 | Datum 1049 | cosine_distance(PG_FUNCTION_ARGS) 1050 | { 1051 | ArrayType *a = PG_GETARG_ARRAYTYPE_P(0); 1052 | ArrayType *b = PG_GETARG_ARRAYTYPE_P(1); 1053 | PG_RETURN_FLOAT4(calc_distance(DIST_COSINE, a, b)); 1054 | } 1055 | 1056 | Datum 1057 | manhattan_distance(PG_FUNCTION_ARGS) 1058 | { 1059 | ArrayType *a = PG_GETARG_ARRAYTYPE_P(0); 1060 | ArrayType *b = PG_GETARG_ARRAYTYPE_P(1); 1061 | PG_RETURN_FLOAT4(calc_distance(DIST_MANHATTAN, a, b)); 1062 | } 1063 | -------------------------------------------------------------------------------- /embedding.control: -------------------------------------------------------------------------------- 1 | comment = 'Vector similarity search with the HNSW algorithm' 2 | default_version = '0.3.6' 3 | module_pathname = '$libdir/embedding' 4 | relocatable = true 5 | -------------------------------------------------------------------------------- /embedding.h: -------------------------------------------------------------------------------- 1 | // Copyright 2023 Neon Inc. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #pragma once 16 | 17 | typedef float coord_t; 18 | typedef float dist_t; 19 | typedef uint32_t idx_t; 20 | typedef uint64_t label_t; 21 | 22 | typedef enum { 23 | DIST_L2, 24 | DIST_COSINE, 25 | DIST_MANHATTAN 26 | } dist_func_t; 27 | 28 | typedef struct 29 | { 30 | size_t dim; 31 | size_t data_size; 32 | size_t offset_data; 33 | size_t offset_label; 34 | size_t size_data_per_element; 35 | size_t elems_per_page; 36 | size_t M; 37 | size_t maxM; 38 | size_t efConstruction; 39 | size_t efSearch; 40 | idx_t enterpoint_node; 41 | dist_func_t dist_func; 42 | } HnswMetadata; 43 | 44 | extern bool hnsw_is_deleted(label_t label); 45 | 46 | extern bool hnsw_search(HnswMetadata* meta, const coord_t *point, size_t* n_results, label_t** results); 47 | extern bool hnsw_bind_point(HnswMetadata* meta, const coord_t *point, idx_t idx); 48 | extern bool hnsw_begin_read(HnswMetadata* meta, idx_t idx, idx_t** indexes, coord_t** coords, label_t* label); 49 | extern void hnsw_end_read(HnswMetadata* meta); 50 | extern void hnsw_begin_write(HnswMetadata* meta, idx_t idx, idx_t** indexes, coord_t** coords, label_t* label); 51 | extern void hnsw_end_write(HnswMetadata* meta); 52 | 53 | extern void hnsw_prefetch(HnswMetadata* meta, idx_t idx); 54 | 55 | extern dist_t hnsw_dist_func(dist_func_t dist, coord_t const* ax, coord_t const* bx, size_t dim); 56 | extern void hnsw_init_dist_func(void); 57 | -------------------------------------------------------------------------------- /hnswalg.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Dmitry Baranchuk 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in all 11 | // copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | // SOFTWARE. 20 | 21 | #include 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | 32 | extern "C" { 33 | #include "embedding.h" 34 | } 35 | 36 | inline dist_t 37 | calc_dist_func(HnswMetadata* meta, coord_t const* ax, coord_t const* bx) 38 | { 39 | return hnsw_dist_func(meta->dist_func, ax, bx, meta->dim); 40 | } 41 | 42 | static std::priority_queue> 43 | searchBaseLayer(HnswMetadata* meta, const coord_t *point, size_t ef) 44 | { 45 | std::vector visited; 46 | const size_t init_visited_size = 64*1024; 47 | coord_t* p_coords; 48 | idx_t* p_indexes; 49 | 50 | visited.resize(init_visited_size); 51 | 52 | std::priority_queue> topResults; 53 | std::priority_queue> candidateSet; 54 | 55 | idx_t enterpoint_node = meta->enterpoint_node; 56 | if (!hnsw_begin_read(meta, enterpoint_node, NULL, &p_coords, NULL)) 57 | return topResults; 58 | 59 | dist_t dist = calc_dist_func(meta, point, p_coords); 60 | hnsw_end_read(meta); 61 | 62 | topResults.emplace(dist, enterpoint_node); 63 | candidateSet.emplace(-dist, enterpoint_node); 64 | visited[enterpoint_node >> 5] = 1 << (enterpoint_node & 31); 65 | dist_t lowerBound = dist; 66 | 67 | while (!candidateSet.empty()) 68 | { 69 | std::pair curr_el_pair = candidateSet.top(); 70 | if (-curr_el_pair.first > lowerBound) 71 | break; 72 | 73 | candidateSet.pop(); 74 | idx_t curNodeNum = curr_el_pair.second; 75 | 76 | hnsw_begin_read(meta, curNodeNum, &p_indexes, NULL, NULL); 77 | size_t size = p_indexes[0]; 78 | 79 | for (size_t j = 0; j < size; ++j) { 80 | size_t tnum = p_indexes[1 + j]; 81 | 82 | if (visited.size() <= (tnum >> 5)) 83 | visited.resize((tnum >> 5) + 1); 84 | 85 | if (!(visited[tnum >> 5] & (1 << (tnum & 31)))) { 86 | hnsw_prefetch(meta, tnum); 87 | } 88 | } 89 | for (size_t j = 0; j < size; ++j) { 90 | size_t tnum = p_indexes[1 + j]; 91 | 92 | if (!(visited[tnum >> 5] & (1 << (tnum & 31)))) { 93 | visited[tnum >> 5] |= 1 << (tnum & 31); 94 | 95 | hnsw_begin_read(meta, tnum, NULL, &p_coords, NULL); 96 | dist = calc_dist_func(meta, point, p_coords); 97 | hnsw_end_read(meta); 98 | 99 | if (topResults.top().first > dist || topResults.size() < ef) { 100 | candidateSet.emplace(-dist, tnum); 101 | 102 | topResults.emplace(dist, tnum); 103 | 104 | if (topResults.size() > ef) 105 | topResults.pop(); 106 | 107 | lowerBound = topResults.top().first; 108 | } 109 | } 110 | } 111 | hnsw_end_read(meta); 112 | } 113 | return topResults; 114 | } 115 | 116 | 117 | void getNeighborsByHeuristic(HnswMetadata* meta, std::priority_queue> &topResults, size_t NN) 118 | { 119 | if (topResults.size() < NN) 120 | return; 121 | 122 | std::priority_queue> resultSet; 123 | std::vector> returnlist; 124 | 125 | while (topResults.size() > 0) { 126 | resultSet.emplace(-topResults.top().first, topResults.top().second); 127 | topResults.pop(); 128 | } 129 | 130 | while (resultSet.size()) { 131 | if (returnlist.size() >= NN) 132 | break; 133 | std::pair curen = resultSet.top(); 134 | dist_t dist_to_query = -curen.first; 135 | resultSet.pop(); 136 | bool good = true; 137 | for (std::pair curen2 : returnlist) { 138 | coord_t *p_coords, *p_coords2; 139 | hnsw_begin_read(meta, curen2.second, NULL, &p_coords2, NULL); 140 | hnsw_begin_read(meta, curen.second, NULL, &p_coords, NULL); 141 | dist_t curdist = calc_dist_func(meta, p_coords2, p_coords); 142 | hnsw_end_read(meta); 143 | hnsw_end_read(meta); 144 | if (curdist < dist_to_query) { 145 | good = false; 146 | break; 147 | } 148 | } 149 | if (good) returnlist.push_back(curen); 150 | } 151 | for (std::pair elem : returnlist) 152 | topResults.emplace(-elem.first, elem.second); 153 | } 154 | 155 | void mutuallyConnectNewElement(HnswMetadata* meta, const coord_t *point, idx_t cur_c, 156 | std::priority_queue> topResults) 157 | { 158 | getNeighborsByHeuristic(meta, topResults, meta->M); 159 | 160 | idx_t *p_indexes; 161 | coord_t *p_coord, *p_coord2; 162 | std::vector res; 163 | res.reserve(meta->M); 164 | while (topResults.size() > 0) { 165 | res.push_back(topResults.top().second); 166 | topResults.pop(); 167 | } 168 | { 169 | hnsw_begin_write(meta, cur_c, &p_indexes, NULL, NULL); 170 | if (*p_indexes) 171 | throw std::runtime_error("Should be blank"); 172 | 173 | *p_indexes++ = res.size(); 174 | 175 | for (size_t idx = 0; idx < res.size(); idx++) { 176 | if (p_indexes[idx]) 177 | throw std::runtime_error("Should be blank"); 178 | p_indexes[idx] = res[idx]; 179 | } 180 | hnsw_end_write(meta); 181 | } 182 | for (size_t idx = 0; idx < res.size(); idx++) { 183 | if (res[idx] == cur_c) 184 | throw std::runtime_error("Connection to the same element"); 185 | 186 | size_t resMmax = meta->maxM; 187 | hnsw_begin_write(meta, res[idx], &p_indexes, &p_coord, NULL); 188 | idx_t sz_link_list_other = *p_indexes; 189 | 190 | if (sz_link_list_other > resMmax || sz_link_list_other < 0) 191 | throw std::runtime_error("Bad sz_link_list_other"); 192 | 193 | if (sz_link_list_other < resMmax) { 194 | p_indexes[1 + sz_link_list_other] = cur_c; 195 | *p_indexes = sz_link_list_other + 1; 196 | } else { 197 | // finding the "weakest" element to replace it with the new one 198 | hnsw_begin_read(meta, cur_c, NULL, &p_coord2, NULL); 199 | dist_t d_max = calc_dist_func(meta, p_coord2, p_coord); 200 | hnsw_end_read(meta); 201 | // Heuristic: 202 | std::priority_queue> candidates; 203 | candidates.emplace(d_max, cur_c); 204 | 205 | for (size_t j = 0; j < sz_link_list_other; j++) 206 | { 207 | hnsw_begin_read(meta, p_indexes[1 + j], NULL, &p_coord2, NULL); 208 | candidates.emplace(calc_dist_func(meta, p_coord2, p_coord), p_indexes[1 + j]); 209 | hnsw_end_read(meta); 210 | } 211 | getNeighborsByHeuristic(meta, candidates, resMmax); 212 | 213 | size_t indx = 0; 214 | while (!candidates.empty()) { 215 | p_indexes[1 + indx] = candidates.top().second; 216 | candidates.pop(); 217 | indx++; 218 | } 219 | *p_indexes = indx; 220 | } 221 | hnsw_end_write(meta); 222 | } 223 | } 224 | 225 | void bindPoint(HnswMetadata* meta, coord_t const* point, idx_t cur_c) 226 | { 227 | // Do nothing for the first element 228 | if (cur_c != 0) { 229 | std::priority_queue > topResults = searchBaseLayer(meta, point, meta->efConstruction); 230 | mutuallyConnectNewElement(meta, point, cur_c, topResults); 231 | } 232 | } 233 | 234 | std::priority_queue> searchKnn(HnswMetadata* meta, const coord_t *query, size_t k) 235 | { 236 | std::priority_queue> topResults; 237 | auto topCandidates = searchBaseLayer(meta, query, k); 238 | while (topCandidates.size() > k) { 239 | topCandidates.pop(); 240 | } 241 | while (!topCandidates.empty()) { 242 | std::pair rez = topCandidates.top(); 243 | label_t label; 244 | hnsw_begin_read(meta, rez.second, NULL, NULL, &label); 245 | if (!hnsw_is_deleted(label)) 246 | topResults.push(std::pair(rez.first, label)); 247 | topCandidates.pop(); 248 | hnsw_end_read(meta); 249 | } 250 | 251 | return topResults; 252 | } 253 | 254 | 255 | 256 | bool hnsw_search(HnswMetadata* meta, const coord_t *point, size_t* n_results, label_t** results) 257 | { 258 | try 259 | { 260 | auto result = searchKnn(meta, point, meta->efSearch); 261 | size_t nResults = result.size(); 262 | *results = (label_t*)malloc(nResults*sizeof(label_t)); 263 | if (*results == NULL) 264 | return false; 265 | for (size_t i = nResults; i-- != 0;) 266 | { 267 | (*results)[i] = result.top().second; 268 | result.pop(); 269 | } 270 | *n_results = nResults; 271 | return true; 272 | } 273 | catch (std::exception& x) 274 | { 275 | return false; 276 | } 277 | } 278 | 279 | bool hnsw_bind_point(HnswMetadata* meta, const coord_t *point, idx_t cur) 280 | { 281 | try 282 | { 283 | bindPoint(meta, point, cur); 284 | return true; 285 | } 286 | catch (std::exception& x) 287 | { 288 | fprintf(stderr, "Catch %s\n", x.what()); 289 | return false; 290 | } 291 | } 292 | 293 | -------------------------------------------------------------------------------- /hnswalg.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Dmitry Baranchuk 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in all 11 | // copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | // SOFTWARE. 20 | 21 | #pragma once 22 | 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | #include 32 | #include 33 | 34 | extern "C" { 35 | #include "embedding.h" 36 | } 37 | 38 | struct HierarchicalNSW : HnswMetadata 39 | { 40 | public: 41 | HierarchicalNSW(size_t dim, size_t maxelements, size_t M, size_t maxM, size_t efConstruction); 42 | ~HierarchicalNSW(); 43 | 44 | 45 | inline coord_t *getDataByInternalId(idx_t internal_id) const { 46 | return (coord_t *)&data_level0_memory[internal_id * size_data_per_element + offset_data]; 47 | } 48 | 49 | inline idx_t *get_linklist0(idx_t internal_id) const { 50 | return (idx_t*)&data_level0_memory[internal_id * size_data_per_element]; 51 | } 52 | 53 | inline label_t *getExternalLabel(idx_t internal_id) const { 54 | return (label_t *)&data_level0_memory[internal_id * size_data_per_element + offset_label]; 55 | } 56 | 57 | std::priority_queue> searchBaseLayer(const coord_t *x, size_t ef); 58 | 59 | void getNeighborsByHeuristic(std::priority_queue> &topResults, size_t NN); 60 | 61 | void mutuallyConnectNewElement(const coord_t *x, idx_t id, std::priority_queue> topResults); 62 | 63 | void addPoint(const coord_t *point, label_t label); 64 | 65 | std::priority_queue> searchKnn(const coord_t *query_data, size_t k); 66 | }; 67 | -------------------------------------------------------------------------------- /test/expected/gh-2.out: -------------------------------------------------------------------------------- 1 | -- https://github.com/neondatabase/pg_embedding/issues/2 2 | SET enable_seqscan = off; 3 | CREATE TABLE t (val real[]); 4 | CREATE INDEX ON t USING hnsw (val) WITH (dims=3, m=3); 5 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 6 | val 7 | ----- 8 | (0 rows) 9 | 10 | DROP TABLE t; 11 | -------------------------------------------------------------------------------- /test/expected/gh-3.out: -------------------------------------------------------------------------------- 1 | -- https://github.com/neondatabase/pg_embedding/issues/3 2 | SET enable_seqscan = off; 3 | CREATE TABLE t(id SERIAL PRIMARY KEY, val REAL[]); 4 | CREATE INDEX ON t using hnsw(val) WITH (dims=3, m=3); 5 | INSERT INTO t (val) VALUES ('{0,1,2}'), ('{1,2,3}'), ('{1,1,1}'); 6 | TRUNCATE t; 7 | INSERT INTO t (val) VALUES ('{4,5,6}'), ('{1,2,3}'), ('{7,8,9}'); 8 | SELECT ctid, id from t order by val <-> ARRAY[3,3,3]; 9 | ctid | id 10 | -------+---- 11 | (0,2) | 5 12 | (0,1) | 4 13 | (0,3) | 6 14 | (3 rows) 15 | 16 | DROP TABLE t; 17 | -------------------------------------------------------------------------------- /test/expected/knn.out: -------------------------------------------------------------------------------- 1 | SET enable_seqscan = off; 2 | CREATE TABLE t (val real[]); 3 | INSERT INTO t (val) VALUES ('{0,1,2}'), ('{1,2,3}'), ('{1,1,1}'), (NULL); 4 | CREATE INDEX ON t USING hnsw (val) WITH (dims=3, m=3); 5 | INSERT INTO t (val) VALUES (array[1,2,4]); 6 | explain SELECT * FROM t ORDER BY val <-> array[3,3,3]; 7 | QUERY PLAN 8 | ------------------------------------------------------------------------ 9 | Index Scan using t_val_idx on t (cost=256.00..260.65 rows=3 width=36) 10 | Order By: (val <-> '{3,3,3}'::real[]) 11 | (2 rows) 12 | 13 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 14 | val 15 | --------- 16 | {1,2,3} 17 | {1,2,4} 18 | {1,1,1} 19 | {0,1,2} 20 | (4 rows) 21 | 22 | SELECT COUNT(*) FROM t; 23 | count 24 | ------- 25 | 5 26 | (1 row) 27 | 28 | CREATE INDEX ON t USING hnsw (val ann_cos_ops) WITH (dims=3, m=3); 29 | explain SELECT * FROM t ORDER BY val <=> array[3,3,3]; 30 | QUERY PLAN 31 | ------------------------------------------------------------------------- 32 | Index Scan using t_val_idx1 on t (cost=256.00..260.65 rows=4 width=36) 33 | Order By: (val <=> '{3,3,3}'::real[]) 34 | (2 rows) 35 | 36 | SELECT * FROM t ORDER BY val <=> array[3,3,3]; 37 | val 38 | --------- 39 | {1,1,1} 40 | {1,2,3} 41 | {1,2,4} 42 | {0,1,2} 43 | (4 rows) 44 | 45 | CREATE INDEX ON t USING hnsw (val ann_manhattan_ops) WITH (dims=3, m=3); 46 | explain SELECT * FROM t ORDER BY val <~> array[3,3,3]; 47 | QUERY PLAN 48 | ------------------------------------------------------------------------- 49 | Index Scan using t_val_idx2 on t (cost=256.00..260.65 rows=4 width=36) 50 | Order By: (val <~> '{3,3,3}'::real[]) 51 | (2 rows) 52 | 53 | SELECT * FROM t ORDER BY val <~> array[3,3,3]; 54 | val 55 | --------- 56 | {1,2,3} 57 | {1,2,4} 58 | {0,1,2} 59 | {1,1,1} 60 | (4 rows) 61 | 62 | SET enable_seqscan = on; 63 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 64 | val 65 | --------- 66 | {1,2,3} 67 | {1,2,4} 68 | {1,1,1} 69 | {0,1,2} 70 | 71 | (5 rows) 72 | 73 | SELECT * FROM t ORDER BY val <=> array[3,3,3]; 74 | val 75 | --------- 76 | {1,1,1} 77 | {1,2,3} 78 | {1,2,4} 79 | {0,1,2} 80 | 81 | (5 rows) 82 | 83 | SELECT * FROM t ORDER BY val <~> array[3,3,3]; 84 | val 85 | --------- 86 | {1,2,3} 87 | {1,2,4} 88 | {0,1,2} 89 | {1,1,1} 90 | 91 | (5 rows) 92 | 93 | delete from t; 94 | vacuum t; 95 | INSERT INTO t (val) VALUES ('{0,1,2}'), ('{1,2,3}'), ('{1,1,1}'), (NULL), (array[1,2,4]); 96 | SET enable_seqscan = off; 97 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 98 | val 99 | --------- 100 | {1,2,3} 101 | {1,2,4} 102 | {1,1,1} 103 | {0,1,2} 104 | (4 rows) 105 | 106 | SELECT * FROM t ORDER BY val <=> array[3,3,3]; 107 | val 108 | --------- 109 | {1,1,1} 110 | {1,2,3} 111 | {1,2,4} 112 | {0,1,2} 113 | (4 rows) 114 | 115 | SELECT * FROM t ORDER BY val <~> array[3,3,3]; 116 | val 117 | --------- 118 | {1,2,3} 119 | {1,2,4} 120 | {0,1,2} 121 | {1,1,1} 122 | (4 rows) 123 | 124 | DROP TABLE t; 125 | -------------------------------------------------------------------------------- /test/sql/gh-2.sql: -------------------------------------------------------------------------------- 1 | -- https://github.com/neondatabase/pg_embedding/issues/2 2 | 3 | SET enable_seqscan = off; 4 | 5 | CREATE TABLE t (val real[]); 6 | CREATE INDEX ON t USING hnsw (val) WITH (dims=3, m=3); 7 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 8 | 9 | 10 | DROP TABLE t; 11 | -------------------------------------------------------------------------------- /test/sql/gh-3.sql: -------------------------------------------------------------------------------- 1 | -- https://github.com/neondatabase/pg_embedding/issues/3 2 | 3 | SET enable_seqscan = off; 4 | 5 | CREATE TABLE t(id SERIAL PRIMARY KEY, val REAL[]); 6 | CREATE INDEX ON t using hnsw(val) WITH (dims=3, m=3); 7 | INSERT INTO t (val) VALUES ('{0,1,2}'), ('{1,2,3}'), ('{1,1,1}'); 8 | TRUNCATE t; 9 | INSERT INTO t (val) VALUES ('{4,5,6}'), ('{1,2,3}'), ('{7,8,9}'); 10 | SELECT ctid, id from t order by val <-> ARRAY[3,3,3]; 11 | 12 | DROP TABLE t; 13 | -------------------------------------------------------------------------------- /test/sql/knn.sql: -------------------------------------------------------------------------------- 1 | SET enable_seqscan = off; 2 | 3 | CREATE TABLE t (val real[]); 4 | INSERT INTO t (val) VALUES ('{0,1,2}'), ('{1,2,3}'), ('{1,1,1}'), (NULL); 5 | CREATE INDEX ON t USING hnsw (val) WITH (dims=3, m=3); 6 | 7 | INSERT INTO t (val) VALUES (array[1,2,4]); 8 | 9 | explain SELECT * FROM t ORDER BY val <-> array[3,3,3]; 10 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 11 | SELECT COUNT(*) FROM t; 12 | 13 | CREATE INDEX ON t USING hnsw (val ann_cos_ops) WITH (dims=3, m=3); 14 | explain SELECT * FROM t ORDER BY val <=> array[3,3,3]; 15 | SELECT * FROM t ORDER BY val <=> array[3,3,3]; 16 | 17 | CREATE INDEX ON t USING hnsw (val ann_manhattan_ops) WITH (dims=3, m=3); 18 | explain SELECT * FROM t ORDER BY val <~> array[3,3,3]; 19 | SELECT * FROM t ORDER BY val <~> array[3,3,3]; 20 | 21 | SET enable_seqscan = on; 22 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 23 | SELECT * FROM t ORDER BY val <=> array[3,3,3]; 24 | SELECT * FROM t ORDER BY val <~> array[3,3,3]; 25 | 26 | delete from t; 27 | vacuum t; 28 | INSERT INTO t (val) VALUES ('{0,1,2}'), ('{1,2,3}'), ('{1,1,1}'), (NULL), (array[1,2,4]); 29 | SET enable_seqscan = off; 30 | SELECT * FROM t ORDER BY val <-> array[3,3,3]; 31 | SELECT * FROM t ORDER BY val <=> array[3,3,3]; 32 | SELECT * FROM t ORDER BY val <~> array[3,3,3]; 33 | 34 | 35 | 36 | DROP TABLE t; 37 | --------------------------------------------------------------------------------