├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.yaml │ └── feature_request.yaml ├── mergify.yml └── workflows │ └── main.yml ├── .gitignore ├── .gitmodules ├── Cargo.toml ├── LICENSE ├── OWNERS ├── README.md ├── build.rs ├── docker-compose.yml ├── examples └── collection.rs ├── src ├── client.rs ├── collection.rs ├── config.rs ├── data.rs ├── error.rs ├── index │ └── mod.rs ├── lib.rs ├── mutate.rs ├── options.rs ├── partition.rs ├── proto │ ├── milvus.proto.common.rs │ ├── milvus.proto.feder.rs │ ├── milvus.proto.milvus.rs │ ├── milvus.proto.msg.rs │ ├── milvus.proto.schema.rs │ └── mod.rs ├── query.rs ├── schema.rs ├── types.rs ├── utils.rs └── value.rs └── tests ├── alias.rs ├── client.rs ├── collection.rs └── common └── mod.rs /.github/ISSUE_TEMPLATE/bug_report.yaml: -------------------------------------------------------------------------------- 1 | name: 🐞 Bug report 2 | description: Create a report to help us improve Milvus-SDK-Rust 3 | title: "[Bug]: " 4 | body: 5 | - type: markdown 6 | attributes: 7 | value: | 8 | Thanks for taking the time to fill out this bug report! 9 | - type: checkboxes 10 | attributes: 11 | label: Is there an existing issue for this? 12 | description: Please search to see if an issue already exists for the bug you encountered. 13 | options: 14 | - label: I have searched the existing issues 15 | required: true 16 | - type: textarea 17 | attributes: 18 | label: Describe the bug 19 | description: A clear and concise description of what the bug is. 20 | placeholder: | 21 | When I do , happens and I see the error message attached below: 22 | ```...``` 23 | validations: 24 | required: false 25 | - type: textarea 26 | attributes: 27 | label: Expected Behavior 28 | description: A clear and concise description of what you expected to happen. 29 | placeholder: When I do , should happen instead. 30 | validations: 31 | required: false 32 | - type: textarea 33 | attributes: 34 | label: Steps/Code To Reproduce behavior 35 | description: | 36 | Follow this [guide](http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports) to craft a minimal bug report. 37 | This helps us reproduce the issue you're having and resolve the issue more quickly. 38 | render: markdown 39 | validations: 40 | required: false 41 | - type: textarea 42 | attributes: 43 | label: Environment details 44 | description: | 45 | Enter the Environment Details: 46 | value: | 47 | - Hardware/Softward conditions (OS, CPU, GPU, Memory): 48 | - Method of installation (Docker, or from source): 49 | - Milvus version (v0.3.1, or v0.4.0): 50 | - Milvus configuration (Settings you made in `server_config.yaml`): 51 | render: markdown 52 | validations: 53 | required: false 54 | - type: textarea 55 | attributes: 56 | label: Anything else? 57 | description: | 58 | If applicable, add screenshots to help explain your problem. 59 | Links? References? Anything that will give us more context about the issue you are encountering! 60 | validations: 61 | required: false -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.yaml: -------------------------------------------------------------------------------- 1 | name: 🚀 Feature request 2 | description: As a user, I want to request a feature for Milvus-SDK-Rust 3 | title: "[Feature]: " 4 | labels: [kind/feature] 5 | assignees: 6 | - yah01 7 | body: 8 | - type: markdown 9 | attributes: 10 | value: | 11 | Thanks for taking the time to request a feature for Milvus-SDK-Rust! Please fill the form in English! 12 | - type: checkboxes 13 | attributes: 14 | label: Is there an existing issue for this? 15 | description: Please search to see if an issue related to this feature request already exists. 16 | options: 17 | - label: I have searched the existing issues 18 | required: true 19 | - type: textarea 20 | attributes: 21 | label: Is your feature request related to a problem? Please describe. 22 | description: A concise description of the problem you are facing or the motivetion behind this feature request. 23 | placeholder: | 24 | I faced a problem due to which ... 25 | validations: 26 | required: true 27 | - type: textarea 28 | attributes: 29 | label: Describe the solution you'd like. 30 | description: A concise description of the solution for the issue. 31 | validations: 32 | required: false 33 | - type: textarea 34 | attributes: 35 | label: Describe an alternate solution. 36 | description: Is there any other approack to solve the problem? 37 | validations: 38 | required: false 39 | - type: textarea 40 | attributes: 41 | label: Anything else? (Additional Context) 42 | description: | 43 | Links? References? Anything that will give us more context about this! 44 | validations: 45 | required: false -------------------------------------------------------------------------------- /.github/mergify.yml: -------------------------------------------------------------------------------- 1 | pull_request_rules: 2 | - name: Add needs-dco label when DCO check failed 3 | conditions: 4 | - base=main 5 | - -status-success=DCO 6 | actions: 7 | label: 8 | remove: 9 | - dco-passed 10 | add: 11 | - needs-dco 12 | comment: 13 | message: | 14 | @{{author}} Thanks for your contribution. Please submit with DCO, see the contributing guide https://github.com/milvus-io/milvus/blob/master/CONTRIBUTING.md#developer-certificate-of-origin-dco. 15 | - name: Add dco-passed label when DCO check passed 16 | conditions: 17 | - base=main 18 | - status-success=DCO 19 | actions: 20 | label: 21 | remove: 22 | - needs-dco 23 | add: 24 | - dco-passed 25 | 26 | - name: Test passed for code changed 27 | conditions: 28 | - base=main 29 | - "status-success=Unittest AMD64 Ubuntu 18.04" 30 | actions: 31 | label: 32 | add: 33 | - ci-passed 34 | 35 | - name: Test passed for document changed 36 | conditions: 37 | - base=main 38 | - -files~=^[^\.]+$ 39 | - -files~=\.(?!md|png) 40 | actions: 41 | label: 42 | add: 43 | - ci-passed 44 | 45 | - name: Test passed for mergify changed 46 | conditions: 47 | - base=main 48 | - -files~=^(?!\.github\/mergify\.yml).*$ 49 | actions: 50 | label: 51 | add: 52 | - ci-passed 53 | 54 | - name: Remove ci-passed when unittest failed 55 | conditions: 56 | - base=main 57 | - "check-failure=Unittest AMD64 Ubuntu 18.04" 58 | - files~=^(?!\.github\/mergify\.yml).*$ 59 | - files~=\.(?!md|png) 60 | actions: 61 | label: 62 | remove: 63 | - ci-passed 64 | 65 | - name: Blocking PR if missing a related issue or PR doesn't have kind/improvement label 66 | conditions: 67 | - base=main 68 | - -body~=\#[0-9]{1,6}(\s+|$) 69 | - -body~=https://github.com/milvus-io/milvus-sdk-rust/issues/[0-9]{1,6}(\s+|$) 70 | - -label=kind/improvement 71 | - -title~=\[automated\] 72 | actions: 73 | label: 74 | add: 75 | - do-not-merge/missing-related-issue 76 | comment: 77 | message: | 78 | @{{author}} Please associate the related issue to the body of your Pull Request. (eg. “issue: #187”) 79 | 80 | 81 | - name: Dismiss block label if related issue be added into PR 82 | conditions: 83 | - or: 84 | - and: 85 | - base=main 86 | - or: 87 | - body~=\#[0-9]{1,6}(\s+|$) 88 | - body~=https://github.com/milvus-io/milvus-sdk-rust/issues/[0-9]{1,6}(\s+|$) 89 | - and: 90 | - base=main 91 | - label=kind/improvement 92 | actions: 93 | label: 94 | remove: 95 | - do-not-merge/missing-related-issue 96 | -------------------------------------------------------------------------------- /.github/workflows/main.yml: -------------------------------------------------------------------------------- 1 | name: Unittest 2 | 3 | on: 4 | push: 5 | paths: 6 | - 'src/**' 7 | - 'tests/**' 8 | - 'milvus-proto/**' 9 | - 'build.rs' 10 | - '.github/**' 11 | - 'docker-compose.yml' 12 | 13 | # Triggers the workflow on push or pull request events but only for the master branch 14 | pull_request: 15 | paths: 16 | - 'src/**' 17 | - 'tests/**' 18 | - 'milvus-proto/**' 19 | - 'build.rs' 20 | - '.github/**' 21 | - 'docker-compose.yml' 22 | 23 | 24 | jobs: 25 | # This workflow contains a single job called "build" 26 | build: 27 | name: Unittest AMD64 Ubuntu ${{ matrix.ubuntu }} 28 | # The type of runner that the job will run on 29 | runs-on: ubuntu-latest 30 | timeout-minutes: 30 31 | strategy: 32 | fail-fast: false 33 | matrix: 34 | ubuntu: [18.04] 35 | env: 36 | UBUNTU: ${{ matrix.ubuntu }} 37 | steps: 38 | # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it 39 | - uses: actions/checkout@v3 40 | with: 41 | submodules: 'true' 42 | 43 | - name: Setup Milvus 44 | run: sudo docker-compose up -d && sleep 15s 45 | 46 | - name: Setup protoc 47 | uses: arduino/setup-protoc@v1.1.2 48 | with: 49 | repo-token: ${{ secrets.GITHUB_TOKEN }} 50 | 51 | # Runs a single command using the runners shell 52 | - name: Run Unittest 53 | run: RUST_BACKTRACE=1 cargo test 54 | # - name: Upload coverage to Codecov 55 | # if: github.repository == 'milvus-io/milvus-sdk-rust' 56 | # uses: codecov/codecov-action@v1 57 | # with: 58 | # token: ${{ secrets.CODECOV_TOKEN }} 59 | # file: ./coverage.project.out 60 | # name: ubuntu-${{ matrix.ubuntu }}-unittests 61 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | /target/ 4 | 5 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 6 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 7 | Cargo.lock 8 | 9 | # These are backup files generated by rustfmt 10 | **/*.rs.bk 11 | 12 | 13 | # Added by cargo 14 | # 15 | # already existing elements were commented out 16 | 17 | /target 18 | #Cargo.lock 19 | /.idea/ 20 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "milvus-proto"] 2 | path = milvus-proto 3 | url = https://github.com/milvus-io/milvus-proto.git 4 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "milvus-sdk-rust" 3 | description = "The official Milvus Rust SDK" 4 | version = "0.1.0" 5 | edition = "2021" 6 | license-file = "LICENSE" 7 | 8 | [lib] 9 | name = "milvus" 10 | 11 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 12 | 13 | [dependencies] 14 | tonic = { version = "0.8.2", features = ["tls", "tls-roots"] } 15 | prost = "0.11.0" 16 | tokio = { version = "1.17.0", features = ["full"] } 17 | thiserror = "1.0" 18 | serde = { version = "1.0", features = ["derive"] } 19 | serde_json = "1.0" 20 | anyhow = "1.0" 21 | strum = "0.24" 22 | strum_macros = "0.24" 23 | base64 = "0.21.0" 24 | dashmap = "5.5.3" 25 | 26 | [build-dependencies] 27 | tonic-build = { version = "0.8.2", default-features = false, features = [ 28 | "prost", 29 | ] } 30 | 31 | [dev-dependencies] 32 | rand = "0.8.5" 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /OWNERS: -------------------------------------------------------------------------------- 1 | filters: 2 | ".*": 3 | reviewers: 4 | - congqixia 5 | - yah01 6 | approvers: 7 | - congqixia 8 | - yah01 9 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Milvus Rust SDK 2 | Rust SDK for Milvus. 3 | 4 | **This is still in progress, be careful to use it in your production, and we are looking for active maintianers of this repo** 5 | 6 | ## Get Started 7 | Add the SDK into your project: 8 | ``` 9 | cargo add milvus-sdk-rust 10 | ``` 11 | 12 | Connect to milvus service and create collection: 13 | ```rust 14 | #[tokio::main] 15 | async fn main() -> Result<(), Error> { 16 | const URL: &str = "http://localhost:19530"; 17 | 18 | let client = Client::new(URL).await?; 19 | 20 | let schema = 21 | CollectionSchemaBuilder::new("hello_milvus", "a guide example for milvus rust SDK") 22 | .add_field(FieldSchema::new_primary_int64( 23 | "id", 24 | "primary key field", 25 | true, 26 | )) 27 | .add_field(FieldSchema::new_float_vector( 28 | DEFAULT_VEC_FIELD, 29 | "feature field", 30 | 256, 31 | )) 32 | .build()?; 33 | let collection = client.create_collection(schema.clone(), None).await?; 34 | Ok(()) 35 | } 36 | ``` 37 | 38 | ## Development 39 | 40 | Pre-requisites: 41 | - cargo 42 | - protocol-compiler 43 | - docker (for testing) 44 | 45 | ### How to test 46 | Many tests require the Milvus server, the project provide a docker-compose file to setup a Milvus cluster: 47 | ``` 48 | docker-compose -f ./docker-compose.yml up -d 49 | ``` 50 | You may need to wait for seconds until the system ready 51 | 52 | Run all tests: 53 | ``` 54 | cargo test 55 | ``` 56 | 57 | Enable the full backtrace for debugging: 58 | ``` 59 | RUST_BACKTRACE=1 cargo test 60 | ``` 61 | -------------------------------------------------------------------------------- /build.rs: -------------------------------------------------------------------------------- 1 | fn main() -> Result<(), Box> { 2 | tonic_build::configure() 3 | .build_server(false) 4 | .out_dir("src/proto") 5 | .compile( 6 | &[ 7 | "milvus-proto/proto/common.proto", 8 | "milvus-proto/proto/milvus.proto", 9 | "milvus-proto/proto/schema.proto", 10 | ], 11 | &["milvus-proto/proto"], 12 | )?; 13 | 14 | Ok(()) 15 | } 16 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '3.5' 2 | 3 | services: 4 | etcd: 5 | container_name: milvus-etcd 6 | image: quay.io/coreos/etcd:v3.5.5 7 | environment: 8 | - ETCD_AUTO_COMPACTION_MODE=revision 9 | - ETCD_AUTO_COMPACTION_RETENTION=1000 10 | - ETCD_QUOTA_BACKEND_BYTES=4294967296 11 | - ETCD_SNAPSHOT_COUNT=50000 12 | volumes: 13 | - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd 14 | command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd 15 | healthcheck: 16 | test: [ "CMD", "curl", "-f", "http://localhost:2379/health" ] 17 | interval: 30s 18 | timeout: 20s 19 | retries: 3 20 | 21 | minio: 22 | container_name: milvus-minio 23 | image: minio/minio:RELEASE.2023-03-20T20-16-18Z 24 | environment: 25 | MINIO_ACCESS_KEY: minioadmin 26 | MINIO_SECRET_KEY: minioadmin 27 | ports: 28 | - "9001:9001" 29 | - "9000:9000" 30 | volumes: 31 | - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data 32 | command: minio server /minio_data --console-address ":9001" 33 | healthcheck: 34 | test: 35 | [ 36 | "CMD", 37 | "curl", 38 | "-f", 39 | "http://localhost:9000/minio/health/live" 40 | ] 41 | interval: 30s 42 | timeout: 20s 43 | retries: 3 44 | 45 | standalone: 46 | container_name: milvus-standalone 47 | image: milvusdb/milvus:master-latest 48 | command: [ "milvus", "run", "standalone" ] 49 | security_opt: 50 | - seccomp:unconfined 51 | environment: 52 | ETCD_ENDPOINTS: etcd:2379 53 | MINIO_ADDRESS: minio:9000 54 | volumes: 55 | - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus 56 | healthcheck: 57 | test: [ "CMD", "curl", "-f", "http://localhost:9091/healthz" ] 58 | interval: 30s 59 | start_period: 90s 60 | timeout: 20s 61 | retries: 3 62 | ports: 63 | - "19530:19530" 64 | - "9091:9091" 65 | depends_on: 66 | - "etcd" 67 | - "minio" 68 | 69 | networks: 70 | default: 71 | name: milvus 72 | -------------------------------------------------------------------------------- /examples/collection.rs: -------------------------------------------------------------------------------- 1 | use milvus::index::{IndexParams, IndexType}; 2 | use milvus::options::LoadOptions; 3 | use milvus::query::QueryOptions; 4 | use milvus::schema::{CollectionSchema, CollectionSchemaBuilder}; 5 | use milvus::{ 6 | client::Client, collection::Collection, data::FieldColumn, error::Error, schema::FieldSchema, 7 | }; 8 | use std::collections::HashMap; 9 | 10 | use rand::prelude::*; 11 | 12 | const DEFAULT_VEC_FIELD: &str = "embed"; 13 | const DIM: i64 = 256; 14 | 15 | #[tokio::main] 16 | async fn main() -> Result<(), Error> { 17 | const URL: &str = "http://localhost:19530"; 18 | 19 | let client = Client::new(URL).await?; 20 | 21 | let schema = 22 | CollectionSchemaBuilder::new("hello_milvus", "a guide example for milvus rust SDK") 23 | .add_field(FieldSchema::new_primary_int64( 24 | "id", 25 | "primary key field", 26 | true, 27 | )) 28 | .add_field(FieldSchema::new_float_vector( 29 | DEFAULT_VEC_FIELD, 30 | "feature field", 31 | DIM, 32 | )) 33 | .build()?; 34 | client.create_collection(schema.clone(), None).await?; 35 | 36 | if let Err(err) = hello_milvus(&client, &schema).await { 37 | println!("failed to run hello milvus: {:?}", err); 38 | } 39 | client.drop_collection(schema.name()).await?; 40 | 41 | Ok(()) 42 | } 43 | 44 | async fn hello_milvus(client: &Client, collection: &CollectionSchema) -> Result<(), Error> { 45 | let mut embed_data = Vec::::new(); 46 | for _ in 1..=DIM * 1000 { 47 | let mut rng = rand::thread_rng(); 48 | let embed = rng.gen(); 49 | embed_data.push(embed); 50 | } 51 | let embed_column = 52 | FieldColumn::new(collection.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); 53 | 54 | client 55 | .insert(collection.name(), vec![embed_column], None) 56 | .await?; 57 | client.flush(collection.name()).await?; 58 | let index_params = IndexParams::new( 59 | "feature_index".to_owned(), 60 | IndexType::IvfFlat, 61 | milvus::index::MetricType::L2, 62 | HashMap::from([("nlist".to_owned(), "32".to_owned())]), 63 | ); 64 | client 65 | .create_index(collection.name(), DEFAULT_VEC_FIELD, index_params) 66 | .await?; 67 | client 68 | .load_collection(collection.name(), Some(LoadOptions::default())) 69 | .await?; 70 | 71 | let options = QueryOptions::default(); 72 | let result = client.query(collection.name(), "id > 0", &options).await?; 73 | 74 | println!( 75 | "result num: {}", 76 | result.first().map(|c| c.len()).unwrap_or(0), 77 | ); 78 | 79 | Ok(()) 80 | } 81 | -------------------------------------------------------------------------------- /src/client.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use crate::collection::CollectionCache; 18 | use crate::config::RPC_TIMEOUT; 19 | use crate::error::{Error, Result}; 20 | pub use crate::proto::common::ConsistencyLevel; 21 | use crate::proto::common::{MsgBase, MsgType}; 22 | use crate::proto::milvus::milvus_service_client::MilvusServiceClient; 23 | use crate::proto::milvus::FlushRequest; 24 | use crate::utils::status_to_result; 25 | use base64::engine::general_purpose; 26 | use base64::Engine; 27 | use std::collections::HashMap; 28 | use std::convert::TryInto; 29 | use std::time::Duration; 30 | use tonic::codegen::{InterceptedService, StdError}; 31 | use tonic::service::Interceptor; 32 | use tonic::transport::Channel; 33 | use tonic::Request; 34 | 35 | #[derive(Clone)] 36 | pub struct AuthInterceptor { 37 | token: Option, 38 | } 39 | 40 | impl Interceptor for AuthInterceptor { 41 | fn call( 42 | &mut self, 43 | mut req: Request<()>, 44 | ) -> std::result::Result, tonic::Status> { 45 | if let Some(ref token) = self.token { 46 | let header_value = format!("{}", token); 47 | req.metadata_mut() 48 | .insert("authorization", header_value.parse().unwrap()); 49 | } 50 | 51 | Ok(req) 52 | } 53 | } 54 | 55 | #[derive(Clone)] 56 | pub struct ClientBuilder { 57 | dst: D, 58 | username: Option, 59 | password: Option, 60 | } 61 | 62 | impl ClientBuilder 63 | where 64 | D: TryInto + Clone, 65 | D::Error: Into, 66 | D::Error: std::fmt::Debug, 67 | { 68 | pub fn new(dst: D) -> Self { 69 | Self { 70 | dst, 71 | username: None, 72 | password: None, 73 | } 74 | } 75 | 76 | pub fn username(mut self, username: &str) -> Self { 77 | self.username = Some(username.to_owned()); 78 | self 79 | } 80 | 81 | pub fn password(mut self, password: &str) -> Self { 82 | self.password = Some(password.to_owned()); 83 | self 84 | } 85 | 86 | pub async fn build(self) -> Result { 87 | Client::with_timeout(self.dst, RPC_TIMEOUT, self.username, self.password).await 88 | } 89 | } 90 | 91 | #[derive(Debug, Clone)] 92 | pub struct Client { 93 | pub(crate) client: MilvusServiceClient>, 94 | pub(crate) collection_cache: CollectionCache, 95 | } 96 | 97 | impl Client { 98 | pub async fn new(dst: D) -> Result 99 | where 100 | D: TryInto, 101 | D::Error: Into + std::fmt::Debug, 102 | { 103 | Self::with_timeout(dst, RPC_TIMEOUT, None, None).await 104 | } 105 | 106 | pub async fn with_timeout( 107 | dst: D, 108 | timeout: Duration, 109 | username: Option, 110 | password: Option, 111 | ) -> Result 112 | where 113 | D: TryInto, 114 | D::Error: Into, 115 | D::Error: std::fmt::Debug, 116 | { 117 | let mut dst: tonic::transport::Endpoint = dst.try_into().map_err(|err| { 118 | Error::InvalidParameter("url".to_owned(), format!("to parse {:?}", err)) 119 | })?; 120 | 121 | dst = dst.timeout(timeout); 122 | 123 | let token = match (username, password) { 124 | (Some(username), Some(password)) => { 125 | let auth_token = format!("{}:{}", username, password); 126 | let auth_token = general_purpose::STANDARD.encode(auth_token); 127 | Some(auth_token) 128 | } 129 | _ => None, 130 | }; 131 | 132 | let auth_interceptor = AuthInterceptor { token }; 133 | 134 | let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; 135 | 136 | let client = MilvusServiceClient::with_interceptor(conn, auth_interceptor); 137 | 138 | Ok(Self { 139 | client: client.clone(), 140 | collection_cache: CollectionCache::new(client), 141 | }) 142 | } 143 | 144 | pub async fn flush_collections(&self, collections: C) -> Result>> 145 | where 146 | C: IntoIterator, 147 | C::Item: ToString, 148 | { 149 | let res = self 150 | .client 151 | .clone() 152 | .flush(FlushRequest { 153 | base: Some(MsgBase::new(MsgType::Flush)), 154 | db_name: "".to_string(), 155 | collection_names: collections.into_iter().map(|x| x.to_string()).collect(), 156 | }) 157 | .await? 158 | .into_inner(); 159 | 160 | status_to_result(&res.status)?; 161 | 162 | Ok(res 163 | .coll_seg_i_ds 164 | .into_iter() 165 | .map(|(k, v)| (k, v.data)) 166 | .collect()) 167 | } 168 | 169 | // alias related: 170 | 171 | /// Creates an alias for a collection. 172 | /// 173 | /// # Arguments 174 | /// 175 | /// * `collection_name` - The name of the collection. 176 | /// * `alias` - The alias to be created. 177 | /// 178 | /// # Returns 179 | /// 180 | /// Returns a `Result` indicating success or failure. 181 | pub async fn create_alias( 182 | &self, 183 | collection_name: impl Into, 184 | alias: impl Into, 185 | ) -> Result<()> { 186 | let collection_name = collection_name.into(); 187 | let alias = alias.into(); 188 | status_to_result(&Some( 189 | self.client 190 | .clone() 191 | .create_alias(crate::proto::milvus::CreateAliasRequest { 192 | base: Some(MsgBase::new(MsgType::CreateAlias)), 193 | db_name: "".to_string(), // reserved 194 | collection_name, 195 | alias, 196 | }) 197 | .await? 198 | .into_inner(), 199 | )) 200 | } 201 | 202 | /// Drops an alias. 203 | /// 204 | /// # Arguments 205 | /// 206 | /// * `alias` - The alias to be dropped. 207 | /// 208 | /// # Returns 209 | /// 210 | /// Returns a `Result` indicating success or failure. 211 | pub async fn drop_alias(&self, alias: S) -> Result<()> 212 | where 213 | S: Into, 214 | { 215 | let alias = alias.into(); 216 | status_to_result(&Some( 217 | self.client 218 | .clone() 219 | .drop_alias(crate::proto::milvus::DropAliasRequest { 220 | base: Some(MsgBase::new(MsgType::DropAlias)), 221 | db_name: "".to_string(), // reserved 222 | alias, 223 | }) 224 | .await? 225 | .into_inner(), 226 | )) 227 | } 228 | 229 | /// Alter the alias of a collection. 230 | /// 231 | /// # Arguments 232 | /// 233 | /// * `collection_name` - The name of the collection. 234 | /// * `alias` - The new alias for the collection. 235 | /// 236 | /// # Returns 237 | /// 238 | /// Returns a `Result` indicating success or failure. 239 | pub async fn alter_alias( 240 | &self, 241 | collection_name: impl Into, 242 | alias: impl Into, 243 | ) -> Result<()> { 244 | let collection_name = collection_name.into(); 245 | let alias = alias.into(); 246 | status_to_result(&Some( 247 | self.client 248 | .clone() 249 | .alter_alias(crate::proto::milvus::AlterAliasRequest { 250 | base: Some(MsgBase::new(MsgType::AlterAlias)), 251 | db_name: "".to_string(), // reserved 252 | collection_name, 253 | alias, 254 | }) 255 | .await? 256 | .into_inner(), 257 | )) 258 | } 259 | } 260 | -------------------------------------------------------------------------------- /src/collection.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use crate::config; 18 | use crate::data::FieldColumn; 19 | use crate::error::{Error as SuperError, Result}; 20 | use crate::index::{IndexInfo, IndexParams}; 21 | use crate::proto::milvus::{ 22 | CreateCollectionRequest, CreateIndexRequest, DescribeIndexRequest, DropCollectionRequest, 23 | DropIndexRequest, FlushRequest, GetCompactionStateRequest, GetCompactionStateResponse, 24 | HasCollectionRequest, LoadCollectionRequest, ManualCompactionRequest, ManualCompactionResponse, 25 | ReleaseCollectionRequest, ShowCollectionsRequest, 26 | }; 27 | use crate::proto::schema::DataType; 28 | use crate::schema::CollectionSchema; 29 | use crate::types::*; 30 | use crate::utils::status_to_result; 31 | use crate::value::Value; 32 | use crate::{ 33 | client::{AuthInterceptor, Client}, 34 | options::{CreateCollectionOptions, GetLoadStateOptions, LoadOptions}, 35 | proto::{ 36 | self, 37 | common::{ConsistencyLevel, IndexState, MsgBase, MsgType}, 38 | milvus::{milvus_service_client::MilvusServiceClient, DescribeCollectionRequest}, 39 | }, 40 | }; 41 | use prost::bytes::BytesMut; 42 | use prost::Message; 43 | use serde_json; 44 | use std::collections::HashMap; 45 | use std::time::Duration; 46 | use thiserror::Error as ThisError; 47 | use tonic::{service::interceptor::InterceptedService, transport::Channel}; 48 | 49 | #[derive(Debug, Clone)] 50 | pub struct Collection { 51 | pub id: i64, 52 | pub name: String, 53 | pub auto_id: bool, 54 | pub num_shards: usize, 55 | // pub num_partitions: usize, 56 | pub consistency_level: ConsistencyLevel, 57 | pub description: String, 58 | pub fields: Vec, 59 | // pub enable_dynamic_field: bool, 60 | } 61 | 62 | #[derive(Debug, Clone)] 63 | pub(crate) struct CollectionCache { 64 | collections: dashmap::DashMap, 65 | timestamps: dashmap::DashMap, 66 | client: MilvusServiceClient>, 67 | } 68 | 69 | impl CollectionCache { 70 | pub fn new(client: MilvusServiceClient>) -> Self { 71 | Self { 72 | collections: dashmap::DashMap::new(), 73 | timestamps: dashmap::DashMap::new(), 74 | client: client, 75 | } 76 | } 77 | 78 | pub async fn get<'a>(&self, name: &str) -> Result { 79 | if !self.local_exist(name) { 80 | let resp = self 81 | .client 82 | .clone() 83 | .describe_collection(DescribeCollectionRequest { 84 | base: Some(MsgBase::new(MsgType::DescribeCollection)), 85 | db_name: "".to_owned(), 86 | collection_name: name.into(), 87 | collection_id: 0, 88 | time_stamp: 0, 89 | }) 90 | .await? 91 | .into_inner(); 92 | 93 | status_to_result(&resp.status)?; 94 | self.collections 95 | .insert(name.to_owned(), Collection::from(resp)); 96 | } 97 | 98 | self.collections 99 | .get(name) 100 | .map(|v| v.value().clone()) 101 | .ok_or(SuperError::Collection(Error::CollectionNotFound( 102 | name.to_owned(), 103 | ))) 104 | } 105 | 106 | pub fn update_timestamp(&self, name: &str, timestamp: Timestamp) { 107 | self.timestamps 108 | .entry(name.to_owned()) 109 | .and_modify(|t| { 110 | if *t < timestamp { 111 | *t = timestamp; 112 | } 113 | }) 114 | .or_insert(timestamp); 115 | } 116 | 117 | pub fn get_timestamp(&self, name: &str) -> Option { 118 | self.timestamps.get(name).map(|v| v.value().clone()) 119 | } 120 | 121 | fn local_exist(&self, name: &str) -> bool { 122 | self.collections.contains_key(name) 123 | } 124 | } 125 | 126 | impl From for Collection { 127 | fn from(value: proto::milvus::DescribeCollectionResponse) -> Self { 128 | let schema = value.schema.unwrap(); 129 | Self { 130 | id: value.collection_id, 131 | name: value.collection_name, 132 | auto_id: schema.auto_id, 133 | num_shards: value.shards_num as usize, 134 | // num_partitions: value.partitions_num as usize, 135 | consistency_level: ConsistencyLevel::from_i32(value.consistency_level).unwrap(), 136 | description: schema.description, 137 | fields: schema.fields.into_iter().map(|f| Field::from(f)).collect(), 138 | // enable_dynamic_field: value.enable_dynamic_field, 139 | } 140 | } 141 | } 142 | 143 | #[derive(Debug)] 144 | pub struct Partition { 145 | pub name: String, 146 | pub percentage: i64, 147 | } 148 | 149 | #[derive(Debug)] 150 | pub struct CompactionInfo { 151 | pub id: i64, 152 | pub plan_count: i32, 153 | } 154 | 155 | impl From for CompactionInfo { 156 | fn from(value: proto::milvus::ManualCompactionResponse) -> Self { 157 | Self { 158 | id: value.compaction_id, 159 | plan_count: value.compaction_plan_count, 160 | } 161 | } 162 | } 163 | 164 | #[derive(Debug)] 165 | pub struct CompactionState { 166 | pub state: crate::proto::common::CompactionState, 167 | pub executing_plan_num: i64, 168 | pub timeout_plan_num: i64, 169 | pub completed_plan_num: i64, 170 | pub failed_plan_num: i64, 171 | } 172 | 173 | impl From for CompactionState { 174 | fn from(value: GetCompactionStateResponse) -> Self { 175 | Self { 176 | state: crate::proto::common::CompactionState::from_i32(value.state).unwrap(), 177 | executing_plan_num: value.executing_plan_no, 178 | timeout_plan_num: value.timeout_plan_no, 179 | completed_plan_num: value.completed_plan_no, 180 | failed_plan_num: value.failed_plan_no, 181 | } 182 | } 183 | } 184 | 185 | type ConcurrentHashMap = tokio::sync::RwLock>; 186 | 187 | impl Client { 188 | /// Creates a new collection with the specified schema and options. 189 | /// 190 | /// # Arguments 191 | /// 192 | /// * `schema` - The schema of the collection. 193 | /// * `options` - Optional parameters for creating the collection. 194 | pub async fn create_collection( 195 | &self, 196 | schema: CollectionSchema, 197 | options: Option, 198 | ) -> Result<()> { 199 | let options = options.unwrap_or_default(); 200 | let schema: crate::proto::schema::CollectionSchema = schema.into(); 201 | let mut buf = BytesMut::new(); 202 | 203 | schema.encode(&mut buf)?; 204 | 205 | let status = self 206 | .client 207 | .clone() 208 | .create_collection(CreateCollectionRequest { 209 | base: Some(MsgBase::new(MsgType::CreateCollection)), 210 | collection_name: schema.name.to_string(), 211 | schema: buf.to_vec(), 212 | shards_num: options.shard_num, 213 | consistency_level: options.consistency_level as i32, 214 | ..Default::default() 215 | }) 216 | .await? 217 | .into_inner(); 218 | 219 | status_to_result(&Some(status)) 220 | } 221 | 222 | /// Drops a collection with the given name. 223 | /// 224 | /// # Arguments 225 | /// 226 | /// * `name` - The name of the collection to drop. 227 | /// 228 | /// # Returns 229 | /// 230 | /// Returns a `Result` indicating success or failure. 231 | pub async fn drop_collection(&self, name: S) -> Result<()> 232 | where 233 | S: Into, 234 | { 235 | status_to_result(&Some( 236 | self.client 237 | .clone() 238 | .drop_collection(DropCollectionRequest { 239 | base: Some(MsgBase::new(MsgType::DropCollection)), 240 | collection_name: name.into(), 241 | ..Default::default() 242 | }) 243 | .await? 244 | .into_inner(), 245 | )) 246 | } 247 | 248 | /// Retrieves a list of collections. 249 | /// 250 | /// # Returns 251 | /// 252 | /// A `Result` containing a vector of collection names if successful, or an error if the operation fails. 253 | pub async fn list_collections(&self) -> Result> { 254 | let response = self 255 | .client 256 | .clone() 257 | .show_collections(ShowCollectionsRequest { 258 | base: Some(MsgBase::new(MsgType::ShowCollections)), 259 | ..Default::default() 260 | }) 261 | .await? 262 | .into_inner(); 263 | 264 | status_to_result(&response.status)?; 265 | Ok(response.collection_names) 266 | } 267 | 268 | /// Retrieves information about a collection. 269 | /// 270 | /// # Arguments 271 | /// 272 | /// * `name` - The name of the collection. 273 | /// 274 | /// # Returns 275 | /// 276 | /// Returns a `Result` containing the `Collection` information if successful, or an error if the collection does not exist or cannot be accessed. 277 | pub async fn describe_collection(&self, name: S) -> Result 278 | where 279 | S: Into, 280 | { 281 | let resp = self 282 | .client 283 | .clone() 284 | .describe_collection(DescribeCollectionRequest { 285 | base: Some(MsgBase::new(MsgType::DescribeCollection)), 286 | db_name: "".to_owned(), 287 | collection_name: name.into(), 288 | collection_id: 0, 289 | time_stamp: 0, 290 | }) 291 | .await? 292 | .into_inner(); 293 | 294 | status_to_result(&resp.status)?; 295 | 296 | Ok(resp.into()) 297 | } 298 | 299 | /// Checks if a collection with the given name exists. 300 | /// 301 | /// # Arguments 302 | /// 303 | /// * `name` - The name of the collection. 304 | /// 305 | /// # Returns 306 | /// 307 | /// Returns a `Result` indicating whether the collection exists or not. 308 | pub async fn has_collection(&self, name: S) -> Result 309 | where 310 | S: Into, 311 | { 312 | let name = name.into(); 313 | let res = self 314 | .client 315 | .clone() 316 | .has_collection(HasCollectionRequest { 317 | base: Some(MsgBase::new(MsgType::HasCollection)), 318 | db_name: "".to_string(), 319 | collection_name: name.clone(), 320 | time_stamp: 0, 321 | }) 322 | .await? 323 | .into_inner(); 324 | 325 | status_to_result(&res.status)?; 326 | 327 | Ok(res.value) 328 | } 329 | 330 | // todo(yah01): implement this after the refactor done 331 | // pub async fn rename_collection(&self, name: S, new_name: S) -> Result<()> 332 | // where 333 | // S: Into, 334 | // { 335 | // let name = name.into(); 336 | // let new_name = new_name.into(); 337 | // let res = self 338 | // .client 339 | // .clone() 340 | // .rename_collection(proto::milvus::RenameCollectionRequest { 341 | // base: Some(MsgBase::new(MsgType::RenameCollection)), 342 | // collection_name: name.clone(), 343 | // new_collection_name: new_name.clone(), 344 | // }) 345 | // .await? 346 | // .into_inner(); 347 | 348 | // status_to_result(&Some(res))?; 349 | 350 | // Ok(()) 351 | // } 352 | 353 | /// Retrieves the statistics of a collection. 354 | /// 355 | /// # Arguments 356 | /// 357 | /// * `name` - The name of the collection. 358 | /// 359 | /// # Returns 360 | /// 361 | /// A `Result` containing a `HashMap` with string keys and string values representing the collection statistics. 362 | pub async fn get_collection_stats(&self, name: &str) -> Result> { 363 | let res = self 364 | .client 365 | .clone() 366 | .get_collection_statistics(proto::milvus::GetCollectionStatisticsRequest { 367 | base: Some(MsgBase::new(MsgType::GetCollectionStatistics)), 368 | db_name: "".into(), 369 | collection_name: name.to_owned(), 370 | }) 371 | .await? 372 | .into_inner(); 373 | 374 | status_to_result(&res.status)?; 375 | 376 | Ok(res.stats.into_iter().map(|s| (s.key, s.value)).collect()) 377 | } 378 | 379 | /// Loads a collection with the given name and options. 380 | /// 381 | /// # Arguments 382 | /// 383 | /// * `collection_name` - The name of the collection to load. 384 | /// * `options` - Optional load options. 385 | /// 386 | /// # Returns 387 | /// 388 | /// Returns a `Result` indicating success or failure. 389 | pub async fn load_collection( 390 | &self, 391 | collection_name: S, 392 | options: Option, 393 | ) -> Result<()> 394 | where 395 | S: Into, 396 | { 397 | let options = options.unwrap_or_default(); 398 | let collection_name = collection_name.into(); 399 | status_to_result(&Some( 400 | self.client 401 | .clone() 402 | .load_collection(LoadCollectionRequest { 403 | base: Some(MsgBase::new(MsgType::LoadCollection)), 404 | db_name: "".to_string(), 405 | collection_name: collection_name.clone(), 406 | replica_number: options.replica_number, 407 | resource_groups: vec![], 408 | refresh: false, 409 | }) 410 | .await? 411 | .into_inner(), 412 | ))?; 413 | 414 | loop { 415 | match self.get_load_state(&collection_name, None).await? { 416 | proto::common::LoadState::NotExist => { 417 | return Err(SuperError::Unexpected("collection not found".to_owned())) 418 | } 419 | proto::common::LoadState::Loading => (), 420 | proto::common::LoadState::Loaded => return Ok(()), 421 | proto::common::LoadState::NotLoad => { 422 | return Err(SuperError::Unexpected("collection not loaded".to_owned())) 423 | } 424 | } 425 | 426 | tokio::time::sleep(Duration::from_millis(config::WAIT_LOAD_DURATION_MS)).await; 427 | } 428 | } 429 | 430 | /// Retrieves the load state of a collection. 431 | /// 432 | /// # Arguments 433 | /// 434 | /// * `collection_name` - The name of the collection. 435 | /// * `options` - Optional parameters for retrieving the load state. 436 | /// 437 | /// # Returns 438 | /// 439 | /// The load state of the collection. 440 | /// 441 | /// # Errors 442 | /// 443 | /// Returns an error if the load state retrieval fails. 444 | pub async fn get_load_state( 445 | &self, 446 | collection_name: S, 447 | options: Option, 448 | ) -> Result 449 | where 450 | S: Into, 451 | { 452 | let options = options.unwrap_or_default(); 453 | let res = self 454 | .client 455 | .clone() 456 | .get_load_state(proto::milvus::GetLoadStateRequest { 457 | base: Some(MsgBase::new(MsgType::Undefined)), 458 | db_name: "".into(), 459 | collection_name: collection_name.into(), 460 | partition_names: options.partition_names, 461 | }) 462 | .await? 463 | .into_inner(); 464 | 465 | status_to_result(&res.status)?; 466 | 467 | Ok(res.state()) 468 | } 469 | 470 | /// Releases a collection with the given name. 471 | /// 472 | /// # Arguments 473 | /// 474 | /// * `collection_name` - The name of the collection to release. 475 | /// 476 | /// # Returns 477 | /// 478 | /// Returns a `Result` indicating success or failure. 479 | pub async fn release_collection(&self, collection_name: S) -> Result<()> 480 | where 481 | S: Into, 482 | { 483 | status_to_result(&Some( 484 | self.client 485 | .clone() 486 | .release_collection(ReleaseCollectionRequest { 487 | base: Some(MsgBase::new(MsgType::ReleaseCollection)), 488 | db_name: "".to_string(), 489 | collection_name: collection_name.into(), 490 | }) 491 | .await? 492 | .into_inner(), 493 | )) 494 | } 495 | 496 | pub async fn flush(&self, collection_name: S) -> Result<()> 497 | where 498 | S: Into, 499 | { 500 | let res = self 501 | .client 502 | .clone() 503 | .flush(FlushRequest { 504 | base: Some(MsgBase::new(MsgType::Flush)), 505 | db_name: "".to_string(), 506 | collection_names: vec![collection_name.into()], 507 | }) 508 | .await? 509 | .into_inner(); 510 | 511 | status_to_result(&res.status)?; 512 | 513 | Ok(()) 514 | } 515 | 516 | async fn create_index_impl( 517 | &self, 518 | collection_name: S, 519 | field_name: impl Into, 520 | index_params: IndexParams, 521 | ) -> Result<()> 522 | where 523 | S: Into, 524 | { 525 | let field_name = field_name.into(); 526 | let status = self 527 | .client 528 | .clone() 529 | .create_index(CreateIndexRequest { 530 | base: Some(MsgBase::new(MsgType::CreateIndex)), 531 | db_name: "".to_string(), 532 | collection_name: collection_name.into(), 533 | field_name, 534 | extra_params: index_params.extra_kv_params(), 535 | index_name: index_params.name().clone(), 536 | }) 537 | .await? 538 | .into_inner(); 539 | status_to_result(&Some(status)) 540 | } 541 | 542 | pub async fn create_index( 543 | &self, 544 | collection_name: S, 545 | field_name: impl Into, 546 | index_params: IndexParams, 547 | ) -> Result<()> 548 | where 549 | S: Into, 550 | { 551 | let collection_name = collection_name.into(); 552 | let field_name = field_name.into(); 553 | self.create_index_impl( 554 | collection_name.clone(), 555 | field_name.clone(), 556 | index_params.clone(), 557 | ) 558 | .await?; 559 | 560 | loop { 561 | let index_infos = self 562 | .describe_index(collection_name.clone(), field_name.clone()) 563 | .await?; 564 | 565 | let index_info = index_infos 566 | .iter() 567 | .find(|&x| x.params().name() == index_params.name()); 568 | if index_info.is_none() { 569 | return Err(SuperError::Unexpected( 570 | "failed to describe index".to_owned(), 571 | )); 572 | } 573 | match index_info.unwrap().state() { 574 | IndexState::Finished => return Ok(()), 575 | IndexState::Failed => return Err(SuperError::Collection(Error::IndexBuildFailed)), 576 | _ => (), 577 | }; 578 | 579 | tokio::time::sleep(Duration::from_millis(config::WAIT_CREATE_INDEX_DURATION_MS)).await; 580 | } 581 | } 582 | 583 | pub async fn describe_index( 584 | &self, 585 | collection_name: S, 586 | field_name: S, 587 | ) -> Result> 588 | where 589 | S: Into, 590 | { 591 | let res = self 592 | .client 593 | .clone() 594 | .describe_index(DescribeIndexRequest { 595 | base: Some(MsgBase::new(MsgType::DescribeIndex)), 596 | db_name: "".to_string(), 597 | collection_name: collection_name.into(), 598 | field_name: field_name.into(), 599 | index_name: "".to_string(), 600 | timestamp: 0, 601 | }) 602 | .await? 603 | .into_inner(); 604 | status_to_result(&res.status)?; 605 | 606 | Ok(res.index_descriptions.into_iter().map(Into::into).collect()) 607 | } 608 | 609 | pub async fn drop_index(&self, collection_name: S, field_name: S) -> Result<()> 610 | where 611 | S: Into, 612 | { 613 | let status = self 614 | .client 615 | .clone() 616 | .drop_index(DropIndexRequest { 617 | base: Some(MsgBase::new(MsgType::DropIndex)), 618 | db_name: "".to_string(), 619 | collection_name: collection_name.into(), 620 | field_name: field_name.into(), 621 | index_name: "".to_string(), 622 | }) 623 | .await? 624 | .into_inner(); 625 | status_to_result(&Some(status)) 626 | } 627 | 628 | pub async fn manual_compaction(&self, collection_name: S) -> Result 629 | where 630 | S: Into, 631 | { 632 | let collection = self.collection_cache.get(&collection_name.into()).await?; 633 | 634 | let resp = self 635 | .client 636 | .clone() 637 | .manual_compaction(ManualCompactionRequest { 638 | collection_id: collection.id, 639 | timetravel: 0, 640 | }) 641 | .await? 642 | .into_inner(); 643 | status_to_result(&resp.status)?; 644 | Ok(resp.into()) 645 | } 646 | 647 | pub async fn get_compaction_state(&self, compaction_id: i64) -> Result { 648 | let resp = self 649 | .client 650 | .clone() 651 | .get_compaction_state(GetCompactionStateRequest { compaction_id }) 652 | .await? 653 | .into_inner(); 654 | status_to_result(&resp.status)?; 655 | Ok(resp.into()) 656 | } 657 | } 658 | 659 | pub type ParamValue = serde_json::Value; 660 | pub use serde_json::json as ParamValue; 661 | 662 | // search result for a single vector 663 | pub struct SearchResult<'a> { 664 | pub size: i64, 665 | pub id: Vec>, 666 | pub field: Vec, 667 | pub score: Vec, 668 | } 669 | 670 | pub struct IndexProgress { 671 | pub total_rows: i64, 672 | pub indexed_rows: i64, 673 | } 674 | 675 | #[derive(Debug, ThisError)] 676 | pub enum Error { 677 | #[error("collection {0} not found")] 678 | CollectionNotFound(String), 679 | 680 | #[error("type mismatched in {0:?}, available types: {1:?}")] 681 | IllegalType(String, Vec), 682 | 683 | #[error("value mismatched in {0:?}, it must be {1:?}")] 684 | IllegalValue(String, String), 685 | 686 | #[error("index build failed")] 687 | IndexBuildFailed, 688 | } 689 | -------------------------------------------------------------------------------- /src/config.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use core::time; 18 | 19 | pub const WAIT_LOAD_DURATION_MS: u64 = 500; 20 | pub const WAIT_CREATE_INDEX_DURATION_MS: u64 = 100; 21 | pub const RPC_TIMEOUT: time::Duration = time::Duration::new(10, 0); 22 | -------------------------------------------------------------------------------- /src/data.rs: -------------------------------------------------------------------------------- 1 | use std::borrow::Cow; 2 | 3 | use crate::{ 4 | proto::schema::{ 5 | self, field_data::Field, scalar_field::Data as ScalarData, 6 | vector_field::Data as VectorData, DataType, ScalarField, VectorField, 7 | }, 8 | schema::FieldSchema, 9 | value::{Value, ValueVec}, 10 | }; 11 | 12 | pub trait HasDataType { 13 | fn data_type() -> DataType; 14 | } 15 | 16 | macro_rules! impl_has_data_type { 17 | ( $($t: ty, $o: expr ),+ ) => {$( 18 | impl HasDataType for $t { 19 | fn data_type() -> DataType { 20 | $o 21 | } 22 | } 23 | )*}; 24 | } 25 | 26 | impl_has_data_type! { 27 | bool, DataType::Bool, 28 | i8, DataType::Int8, 29 | i16, DataType::Int16, 30 | i32, DataType::Int32, 31 | i64, DataType::Int64, 32 | f32, DataType::Float, 33 | f64, DataType::Double, 34 | String, DataType::String, 35 | Cow<'_, str>, DataType::String, 36 | Vec, DataType::FloatVector, 37 | Vec, DataType::BinaryVector, 38 | Cow<'_, [f32]>, DataType::FloatVector, 39 | Cow<'_, [u8]>, DataType::BinaryVector 40 | } 41 | 42 | #[derive(Debug, Clone)] 43 | pub struct FieldColumn { 44 | pub name: String, 45 | pub dtype: DataType, 46 | pub value: ValueVec, 47 | pub dim: i64, 48 | pub max_length: i32, 49 | pub is_dynamic: bool, 50 | } 51 | 52 | impl From for FieldColumn { 53 | fn from(fd: schema::FieldData) -> Self { 54 | let (dim, max_length) = fd 55 | .field 56 | .as_ref() 57 | .map(get_dim_max_length) 58 | .unwrap_or((Some(1), None)); 59 | 60 | let value: ValueVec = fd.field.map(Into::into).unwrap_or(ValueVec::None); 61 | let dtype = DataType::from_i32(fd.r#type).unwrap_or(DataType::None); 62 | 63 | FieldColumn { 64 | name: fd.field_name, 65 | dtype, 66 | dim: dim.unwrap(), 67 | max_length: max_length.unwrap_or(0), 68 | value, 69 | is_dynamic: fd.is_dynamic, 70 | } 71 | } 72 | } 73 | 74 | impl FieldColumn { 75 | pub fn new>(schm: &FieldSchema, v: V) -> FieldColumn { 76 | FieldColumn { 77 | name: schm.name.clone(), 78 | dtype: schm.dtype, 79 | value: v.into(), 80 | dim: schm.dim, 81 | max_length: schm.max_length, 82 | is_dynamic: false, 83 | } 84 | } 85 | 86 | pub fn get(&self, idx: usize) -> Option> { 87 | Some(match &self.value { 88 | ValueVec::None => Value::None, 89 | ValueVec::Bool(v) => Value::Bool(*v.get(idx)?), 90 | ValueVec::Int(v) => match self.dtype { 91 | DataType::Int8 => Value::Int8(*v.get(idx)? as _), 92 | DataType::Int16 => Value::Int16(*v.get(idx)? as _), 93 | DataType::Int32 => Value::Int32(*v.get(idx)?), 94 | _ => unreachable!(), 95 | }, 96 | ValueVec::Long(v) => Value::Long(*v.get(idx)?), 97 | ValueVec::Float(v) => match self.dtype { 98 | DataType::Float => Value::Float(*v.get(idx)?), 99 | DataType::FloatVector => { 100 | let dim = self.dim as usize; 101 | Value::FloatArray(Cow::Borrowed(&v[idx * dim..idx * dim + dim])) 102 | } 103 | _ => unreachable!(), 104 | }, 105 | ValueVec::Double(v) => Value::Double(*v.get(idx)?), 106 | ValueVec::Binary(v) => { 107 | let dim = (self.dim / 8) as usize; 108 | Value::Binary(Cow::Borrowed(&v[idx * dim..idx * dim + dim])) 109 | } 110 | ValueVec::String(v) => Value::String(Cow::Borrowed(v.get(idx)?.as_ref())), 111 | ValueVec::Json(v) => Value::Json(Cow::Borrowed(v.get(idx)?.as_ref())), 112 | ValueVec::Array(v) => Value::Array(Cow::Borrowed(v.get(idx)?)), 113 | }) 114 | } 115 | 116 | pub fn push(&mut self, val: Value) { 117 | match (&mut self.value, val) { 118 | (ValueVec::None, Value::None) => (), 119 | (ValueVec::Bool(vec), Value::Bool(i)) => vec.push(i), 120 | (ValueVec::Int(vec), Value::Int8(i)) => vec.push(i as _), 121 | (ValueVec::Int(vec), Value::Int16(i)) => vec.push(i as _), 122 | (ValueVec::Int(vec), Value::Int32(i)) => vec.push(i), 123 | (ValueVec::Long(vec), Value::Long(i)) => vec.push(i), 124 | (ValueVec::Float(vec), Value::Float(i)) => vec.push(i), 125 | (ValueVec::Double(vec), Value::Double(i)) => vec.push(i), 126 | (ValueVec::String(vec), Value::String(i)) => vec.push(i.to_string()), 127 | (ValueVec::Binary(vec), Value::Binary(i)) => vec.extend_from_slice(i.as_ref()), 128 | (ValueVec::Float(vec), Value::FloatArray(i)) => vec.extend_from_slice(i.as_ref()), 129 | _ => panic!("column type mismatch"), 130 | } 131 | } 132 | 133 | #[inline] 134 | pub fn len(&self) -> usize { 135 | self.value.len() / self.dim as usize 136 | } 137 | 138 | pub fn copy_with_metadata(&self) -> Self { 139 | Self { 140 | dim: self.dim, 141 | dtype: self.dtype, 142 | max_length: self.max_length, 143 | name: self.name.clone(), 144 | value: match &self.value { 145 | ValueVec::None => ValueVec::None, 146 | ValueVec::Bool(_) => ValueVec::Bool(Vec::new()), 147 | ValueVec::Int(_) => ValueVec::Int(Vec::new()), 148 | ValueVec::Long(_) => ValueVec::Long(Vec::new()), 149 | ValueVec::Float(_) => ValueVec::Float(Vec::new()), 150 | ValueVec::Double(_) => ValueVec::Double(Vec::new()), 151 | ValueVec::String(_) => ValueVec::String(Vec::new()), 152 | ValueVec::Json(_) => ValueVec::Json(Vec::new()), 153 | ValueVec::Binary(_) => ValueVec::Binary(Vec::new()), 154 | ValueVec::Array(_) => ValueVec::Array(Vec::new()), 155 | }, 156 | is_dynamic: self.is_dynamic, 157 | } 158 | } 159 | } 160 | 161 | impl From for schema::FieldData { 162 | fn from(this: FieldColumn) -> schema::FieldData { 163 | schema::FieldData { 164 | field_name: this.name.to_string(), 165 | field_id: 0, 166 | r#type: this.dtype as _, 167 | field: Some(match this.value { 168 | ValueVec::None => Field::Scalars(ScalarField { data: None }), 169 | ValueVec::Bool(v) => Field::Scalars(ScalarField { 170 | data: Some(ScalarData::BoolData(schema::BoolArray { data: v })), 171 | }), 172 | ValueVec::Int(v) => Field::Scalars(ScalarField { 173 | data: Some(ScalarData::IntData(schema::IntArray { data: v })), 174 | }), 175 | ValueVec::Long(v) => Field::Scalars(ScalarField { 176 | data: Some(ScalarData::LongData(schema::LongArray { data: v })), 177 | }), 178 | ValueVec::Float(v) => match this.dtype { 179 | DataType::Float => Field::Scalars(ScalarField { 180 | data: Some(ScalarData::FloatData(schema::FloatArray { data: v })), 181 | }), 182 | DataType::FloatVector => Field::Vectors(VectorField { 183 | data: Some(VectorData::FloatVector(schema::FloatArray { data: v })), 184 | dim: this.dim, 185 | }), 186 | _ => unimplemented!(), 187 | }, 188 | ValueVec::Double(v) => Field::Scalars(ScalarField { 189 | data: Some(ScalarData::DoubleData(schema::DoubleArray { data: v })), 190 | }), 191 | ValueVec::String(v) => Field::Scalars(ScalarField { 192 | data: Some(ScalarData::StringData(schema::StringArray { data: v })), 193 | }), 194 | ValueVec::Json(v) => Field::Scalars(ScalarField { 195 | data: Some(ScalarData::JsonData(schema::JsonArray { data: v })), 196 | }), 197 | ValueVec::Array(v) => Field::Scalars(ScalarField { 198 | data: Some(ScalarData::ArrayData(schema::ArrayArray { 199 | data: v, 200 | element_type: this.dtype as _, 201 | })), 202 | }), 203 | ValueVec::Binary(v) => Field::Vectors(VectorField { 204 | data: Some(VectorData::BinaryVector(v)), 205 | dim: this.dim, 206 | }), 207 | }), 208 | is_dynamic: false, 209 | } 210 | } 211 | } 212 | 213 | pub trait FromField: Sized { 214 | fn from_field(field: Field) -> Option; 215 | } 216 | 217 | macro_rules! impl_from_field { 218 | ( $( $t: ty [$($e:tt)*] ),+ ) => {$( 219 | 220 | impl FromField for $t { 221 | fn from_field(v: Field) -> Option { 222 | match v { 223 | $($e)*, 224 | _ => None 225 | } 226 | } 227 | } 228 | )*}; 229 | } 230 | 231 | impl_from_field! { 232 | Vec[Field::Scalars(ScalarField {data: Some(ScalarData::BoolData(schema::BoolArray { data }))}) => Some(data)], 233 | Vec[Field::Scalars(ScalarField {data: Some(ScalarData::IntData(schema::IntArray { data }))}) => Some(data.into_iter().map(|x|x as _).collect())], 234 | Vec[Field::Scalars(ScalarField {data: Some(ScalarData::IntData(schema::IntArray { data }))}) => Some(data.into_iter().map(|x|x as _).collect())], 235 | Vec[Field::Scalars(ScalarField {data: Some(ScalarData::IntData(schema::IntArray { data }))}) => Some(data)], 236 | Vec[Field::Scalars(ScalarField {data: Some(ScalarData::LongData(schema::LongArray { data }))}) => Some(data)], 237 | Vec[Field::Scalars(ScalarField {data: Some(ScalarData::StringData(schema::StringArray { data }))}) => Some(data)], 238 | Vec[Field::Scalars(ScalarField {data: Some(ScalarData::DoubleData(schema::DoubleArray { data }))}) => Some(data)], 239 | Vec[Field::Vectors(VectorField {data: Some(VectorData::BinaryVector(data)), ..}) => Some(data)] 240 | } 241 | 242 | impl FromField for Vec { 243 | fn from_field(field: Field) -> Option { 244 | match field { 245 | Field::Scalars(ScalarField { 246 | data: Some(ScalarData::FloatData(schema::FloatArray { data })), 247 | }) => Some(data), 248 | 249 | Field::Vectors(VectorField { 250 | data: Some(VectorData::FloatVector(schema::FloatArray { data })), 251 | .. 252 | }) => Some(data), 253 | 254 | _ => None, 255 | } 256 | } 257 | } 258 | 259 | fn get_dim_max_length(field: &Field) -> (Option, Option) { 260 | let dim = match field { 261 | Field::Scalars(ScalarField { data: Some(_) }) => 1i64, 262 | Field::Vectors(VectorField { dim, .. }) => *dim, 263 | _ => 0i64, 264 | }; 265 | 266 | (Some(dim), None) // no idea how to get max_length 267 | } 268 | -------------------------------------------------------------------------------- /src/error.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | // For custom error handling 17 | // TODO 18 | 19 | use crate::collection::Error as CollectionError; 20 | use crate::proto::common::{ErrorCode, Status}; 21 | use crate::schema::Error as SchemaError; 22 | use std::result; 23 | use thiserror::Error; 24 | use tonic::transport::Error as CommError; 25 | use tonic::Status as GrpcError; 26 | 27 | #[derive(Debug, Error)] 28 | pub enum Error { 29 | #[error("{0:?}")] 30 | Communication(#[from] CommError), 31 | 32 | #[error("{0:?}")] 33 | Collection(#[from] CollectionError), 34 | 35 | #[error("{0:?}")] 36 | Grpc(#[from] GrpcError), 37 | 38 | #[error("{0:?}")] 39 | Schema(#[from] SchemaError), 40 | 41 | #[error("{0:?} {1:?}")] 42 | Server(ErrorCode, String), 43 | 44 | #[error("{0:?}")] 45 | ProstEncode(#[from] prost::EncodeError), 46 | 47 | #[error("{0:?}")] 48 | ProstDecode(#[from] prost::DecodeError), 49 | 50 | #[error("Conversion error")] 51 | Conversion, 52 | #[error("{0:?}")] 53 | SerdeJsonErr(#[from] serde_json::Error), 54 | 55 | #[error("parameter {0:?} with invalid value {1:?}")] 56 | InvalidParameter(String, String), 57 | 58 | #[error("{0:?}")] 59 | Other(#[from] anyhow::Error), 60 | 61 | #[error("{0}")] 62 | Unexpected(String), 63 | } 64 | 65 | impl From for Error { 66 | fn from(s: Status) -> Self { 67 | Error::Server(ErrorCode::from_i32(s.error_code).unwrap(), s.reason) 68 | } 69 | } 70 | 71 | pub type Result = result::Result; 72 | -------------------------------------------------------------------------------- /src/index/mod.rs: -------------------------------------------------------------------------------- 1 | use strum_macros::{Display, EnumString}; 2 | 3 | use crate::proto::{ 4 | common::{IndexState, KeyValuePair}, 5 | milvus::IndexDescription, 6 | }; 7 | use std::{collections::HashMap, str::FromStr}; 8 | 9 | #[derive(Debug, Clone, Copy, EnumString, Display)] 10 | pub enum IndexType { 11 | #[strum(serialize = "FLAT")] 12 | Flat, 13 | #[strum(serialize = "BIN_FLAT")] 14 | BinFlat, 15 | #[strum(serialize = "IVF_FLAT")] 16 | IvfFlat, 17 | #[strum(serialize = "BIN_IVF_FLAT")] 18 | BinIvfFlat, 19 | #[strum(serialize = "IVF_PQ")] 20 | IvfPQ, 21 | #[strum(serialize = "IVF_SQ8")] 22 | IvfSQ8, 23 | #[strum(serialize = "IVF_SQ8_HYBRID")] 24 | IvfSQ8H, 25 | #[strum(serialize = "NSG")] 26 | NSG, 27 | #[strum(serialize = "HNSW")] 28 | HNSW, 29 | #[strum(serialize = "RHNSW_FLAT")] 30 | RHNSWFlat, 31 | #[strum(serialize = "RHNSW_PQ")] 32 | RHNSWPQ, 33 | #[strum(serialize = "RHNSW_SQ")] 34 | RHNSWSQ, 35 | #[strum(serialize = "IVF_HNSW")] 36 | IvfHNSW, 37 | #[strum(serialize = "ANNOY")] 38 | ANNOY, 39 | #[strum(serialize = "NGT_PANNG")] 40 | NGTPANNG, 41 | #[strum(serialize = "NGT_ONNG")] 42 | NGTONNG, 43 | } 44 | 45 | #[derive(Debug, Clone, Copy, EnumString, Display)] 46 | pub enum MetricType { 47 | L2, 48 | IP, 49 | HAMMING, 50 | JACCARD, 51 | TANIMOTO, 52 | SUBSTRUCTURE, 53 | SUPERSTRUCTURE, 54 | } 55 | 56 | #[derive(Debug, Clone)] 57 | pub struct IndexParams { 58 | name: String, 59 | index_type: IndexType, 60 | metric_type: MetricType, 61 | params: HashMap, 62 | } 63 | 64 | impl IndexParams { 65 | pub fn new( 66 | name: String, 67 | index_type: IndexType, 68 | metric_type: MetricType, 69 | params: HashMap, 70 | ) -> Self { 71 | Self { 72 | name, 73 | index_type, 74 | metric_type, 75 | params, 76 | } 77 | } 78 | 79 | pub fn name(&self) -> &String { 80 | &self.name 81 | } 82 | 83 | pub fn index_type(&self) -> IndexType { 84 | self.index_type 85 | } 86 | 87 | pub fn metric_type(&self) -> MetricType { 88 | self.metric_type 89 | } 90 | 91 | pub fn params(&self) -> &HashMap { 92 | &self.params 93 | } 94 | 95 | pub fn extra_params(&self) -> HashMap { 96 | HashMap::from([ 97 | ("index_type".to_owned(), self.index_type().to_string()), 98 | ("metric_type".to_owned(), self.metric_type().to_string()), 99 | ( 100 | "params".to_owned(), 101 | serde_json::to_string(&self.params()).unwrap(), 102 | ), 103 | ]) 104 | } 105 | 106 | pub fn extra_kv_params(&self) -> Vec { 107 | self.extra_params() 108 | .into_iter() 109 | .map(|(k, v)| KeyValuePair { key: k, value: v }) 110 | .collect() 111 | } 112 | } 113 | 114 | #[derive(Debug, Clone)] 115 | pub struct IndexInfo { 116 | field_name: String, 117 | id: i64, 118 | params: IndexParams, 119 | state: IndexState, 120 | } 121 | 122 | impl IndexInfo { 123 | pub fn field_name(&self) -> &str { 124 | &self.field_name 125 | } 126 | 127 | pub fn id(&self) -> i64 { 128 | self.id 129 | } 130 | 131 | pub fn params(&self) -> &IndexParams { 132 | &self.params 133 | } 134 | 135 | pub fn state(&self) -> IndexState { 136 | self.state 137 | } 138 | } 139 | 140 | impl From for IndexInfo { 141 | fn from(description: IndexDescription) -> Self { 142 | let mut params: HashMap = HashMap::from_iter( 143 | description 144 | .params 145 | .iter() 146 | .map(|kv| (kv.key.clone(), kv.value.clone())), 147 | ); 148 | 149 | let index_type = IndexType::from_str(¶ms.remove("index_type").unwrap()).unwrap(); 150 | let metric_type = MetricType::from_str(¶ms.remove("metric_type").unwrap()).unwrap(); 151 | let params = serde_json::from_str(params.get("params").unwrap()).unwrap(); 152 | 153 | let params = IndexParams::new( 154 | description.index_name.clone(), 155 | index_type, 156 | metric_type, 157 | params, 158 | ); 159 | Self { 160 | field_name: description.field_name.clone(), 161 | id: description.index_id, 162 | params: params, 163 | state: description.state(), 164 | } 165 | } 166 | } 167 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | pub mod client; 18 | pub mod collection; 19 | pub mod data; 20 | pub mod error; 21 | pub mod mutate; 22 | pub mod options; 23 | pub mod partition; 24 | pub mod query; 25 | pub mod schema; 26 | pub mod value; 27 | 28 | mod config; 29 | pub mod index; 30 | pub mod proto; 31 | pub mod types; 32 | mod utils; 33 | -------------------------------------------------------------------------------- /src/mutate.rs: -------------------------------------------------------------------------------- 1 | use prost::bytes::{BufMut, BytesMut}; 2 | 3 | use crate::error::Result; 4 | use crate::{ 5 | client::Client, 6 | collection, 7 | data::FieldColumn, 8 | error::Error, 9 | proto::{ 10 | self, 11 | common::{MsgBase, MsgType}, 12 | milvus::{InsertRequest, UpsertRequest}, 13 | schema::{scalar_field::Data, DataType}, 14 | }, 15 | schema::FieldData, 16 | utils::status_to_result, 17 | value::ValueVec, 18 | }; 19 | 20 | #[derive(Debug, Clone)] 21 | pub struct InsertOptions { 22 | pub(crate) partition_name: String, 23 | } 24 | 25 | impl Default for InsertOptions { 26 | fn default() -> Self { 27 | Self { 28 | partition_name: String::new(), 29 | } 30 | } 31 | } 32 | 33 | impl InsertOptions { 34 | pub fn new() -> Self { 35 | Self::default() 36 | } 37 | 38 | pub fn with_partition_name(partition_name: String) -> Self { 39 | Self::default().partition_name(partition_name) 40 | } 41 | 42 | pub fn partition_name(mut self, partition_name: String) -> Self { 43 | self.partition_name = partition_name.to_owned(); 44 | self 45 | } 46 | } 47 | 48 | #[derive(Debug, Clone)] 49 | pub struct DeleteOptions { 50 | pub(crate) ids: ValueVec, 51 | pub(crate) filter: String, 52 | pub(crate) partition_name: String, 53 | } 54 | 55 | impl DeleteOptions { 56 | fn new() -> Self { 57 | Self { 58 | ids: ValueVec::None, 59 | filter: String::new(), 60 | partition_name: String::new(), 61 | } 62 | } 63 | 64 | pub fn with_ids(ids: ValueVec) -> Self { 65 | let mut opt = Self::new(); 66 | opt.ids = ids; 67 | opt 68 | } 69 | 70 | pub fn with_filter(filter: String) -> Self { 71 | let mut opt = Self::new(); 72 | opt.filter = filter; 73 | opt 74 | } 75 | 76 | pub fn partition_name(mut self, partition_name: String) -> Self { 77 | self.partition_name = partition_name; 78 | self 79 | } 80 | } 81 | 82 | impl Client { 83 | pub async fn insert( 84 | &self, 85 | collection_name: S, 86 | fields_data: Vec, 87 | options: Option, 88 | ) -> Result 89 | where 90 | S: Into, 91 | { 92 | let options = options.unwrap_or_default(); 93 | let row_num = fields_data.first().map(|c| c.len()).unwrap_or(0); 94 | let collection_name = collection_name.into(); 95 | 96 | let result = self 97 | .client 98 | .clone() 99 | .insert(InsertRequest { 100 | base: Some(MsgBase::new(MsgType::Insert)), 101 | db_name: "".to_string(), 102 | collection_name: collection_name.clone(), 103 | partition_name: options.partition_name, 104 | num_rows: row_num as u32, 105 | fields_data: fields_data.into_iter().map(|f| f.into()).collect(), 106 | hash_keys: Vec::new(), 107 | }) 108 | .await? 109 | .into_inner(); 110 | 111 | self.collection_cache 112 | .update_timestamp(&collection_name, result.timestamp); 113 | 114 | Ok(result) 115 | } 116 | 117 | pub async fn delete( 118 | &self, 119 | collection_name: impl Into, 120 | options: &DeleteOptions, 121 | ) -> Result { 122 | let collection_name = collection_name.into(); 123 | 124 | let expr = self.compose_expr(&collection_name, options).await?; 125 | 126 | let result = self 127 | .client 128 | .clone() 129 | .delete(proto::milvus::DeleteRequest { 130 | base: Some(MsgBase::new(MsgType::Delete)), 131 | db_name: "".to_string(), 132 | collection_name: collection_name.clone(), 133 | expr: expr, 134 | partition_name: options.partition_name.clone(), 135 | hash_keys: Vec::new(), 136 | }) 137 | .await? 138 | .into_inner(); 139 | 140 | self.collection_cache 141 | .update_timestamp(&collection_name, result.timestamp); 142 | 143 | Ok(result) 144 | } 145 | 146 | async fn compose_expr(&self, collection_name: &str, options: &DeleteOptions) -> Result { 147 | let mut expr = String::new(); 148 | match options.filter.len() { 149 | 0 => { 150 | let collection = self.collection_cache.get(collection_name).await?; 151 | let pk = collection.fields.iter().find(|f| f.is_primary_key).unwrap(); 152 | 153 | let mut expr = String::new(); 154 | expr.push_str(&pk.name); 155 | expr.push_str(" in ["); 156 | match (pk.dtype, options.ids.clone()) { 157 | (DataType::Int64, ValueVec::Long(values)) => { 158 | for (i, v) in values.iter().enumerate() { 159 | if i > 0 { 160 | expr.push_str(","); 161 | } 162 | expr.push_str(format!("{}", v).as_str()); 163 | } 164 | expr 165 | } 166 | 167 | (DataType::VarChar, ValueVec::String(values)) => { 168 | for (i, v) in values.iter().enumerate() { 169 | if i > 0 { 170 | expr.push_str(","); 171 | } 172 | expr.push_str(v.as_str()); 173 | } 174 | expr 175 | } 176 | 177 | _ => { 178 | return Err(Error::InvalidParameter( 179 | "pk type".to_owned(), 180 | pk.dtype.as_str_name().to_owned(), 181 | )); 182 | } 183 | } 184 | } 185 | 186 | _ => options.filter.clone(), 187 | }; 188 | expr.push(')'); 189 | 190 | Ok(expr) 191 | } 192 | 193 | pub async fn upsert( 194 | &self, 195 | collection_name: S, 196 | fields_data: Vec, 197 | options: Option, 198 | ) -> Result 199 | where 200 | S: Into, 201 | { 202 | let options = options.unwrap_or_default(); 203 | let row_num = fields_data.first().map(|c| c.len()).unwrap_or(0); 204 | let collection_name = collection_name.into(); 205 | 206 | let result = self 207 | .client 208 | .clone() 209 | .upsert(UpsertRequest { 210 | base: Some(MsgBase::new(MsgType::Upsert)), 211 | db_name: "".to_string(), 212 | collection_name: collection_name.clone(), 213 | partition_name: options.partition_name, 214 | num_rows: row_num as u32, 215 | fields_data: fields_data.into_iter().map(|f| f.into()).collect(), 216 | hash_keys: Vec::new(), 217 | }) 218 | .await? 219 | .into_inner(); 220 | 221 | self.collection_cache 222 | .update_timestamp(&collection_name, result.timestamp); 223 | 224 | Ok(result) 225 | } 226 | } 227 | -------------------------------------------------------------------------------- /src/options.rs: -------------------------------------------------------------------------------- 1 | use crate::proto::common::ConsistencyLevel; 2 | 3 | #[derive(Debug, Clone, Copy)] 4 | pub struct CreateCollectionOptions { 5 | pub(crate) shard_num: i32, 6 | pub(crate) consistency_level: ConsistencyLevel, 7 | } 8 | 9 | impl Default for CreateCollectionOptions { 10 | fn default() -> Self { 11 | Self { 12 | shard_num: 0, 13 | consistency_level: ConsistencyLevel::Bounded, 14 | } 15 | } 16 | } 17 | 18 | impl CreateCollectionOptions { 19 | pub fn new() -> Self { 20 | Self::default() 21 | } 22 | 23 | pub fn with_shard_num(shard_num: i32) -> Self { 24 | Self::default().shard_num(shard_num) 25 | } 26 | 27 | pub fn with_consistency_level(consistency_level: ConsistencyLevel) -> Self { 28 | Self::default().consistency_level(consistency_level) 29 | } 30 | 31 | pub fn shard_num(mut self, shard_num: i32) -> Self { 32 | self.shard_num = shard_num; 33 | self 34 | } 35 | 36 | pub fn consistency_level(mut self, consistency_level: ConsistencyLevel) -> Self { 37 | self.consistency_level = consistency_level; 38 | self 39 | } 40 | } 41 | 42 | #[derive(Debug, Clone, Copy)] 43 | pub struct LoadOptions { 44 | pub(crate) replica_number: i32, 45 | } 46 | 47 | impl Default for LoadOptions { 48 | fn default() -> Self { 49 | Self { replica_number: 1 } 50 | } 51 | } 52 | 53 | impl LoadOptions { 54 | pub fn new() -> Self { 55 | Self::default() 56 | } 57 | 58 | pub fn with_replica_number(replica_number: i32) -> Self { 59 | Self::default().replica_number(replica_number) 60 | } 61 | 62 | pub fn replica_number(mut self, replica_number: i32) -> Self { 63 | self.replica_number = replica_number; 64 | self 65 | } 66 | } 67 | 68 | #[derive(Debug, Clone)] 69 | pub struct GetLoadStateOptions { 70 | pub(crate) partition_names: Vec, 71 | } 72 | 73 | impl Default for GetLoadStateOptions { 74 | fn default() -> Self { 75 | Self { 76 | partition_names: vec![], 77 | } 78 | } 79 | } 80 | 81 | impl GetLoadStateOptions { 82 | pub fn new() -> Self { 83 | Self::default() 84 | } 85 | 86 | pub fn with_partition_names(partition_names: Vec) -> Self { 87 | Self::default().partition_names(partition_names) 88 | } 89 | 90 | pub fn partition_names(mut self, partition_names: Vec) -> Self { 91 | self.partition_names = partition_names; 92 | self 93 | } 94 | } 95 | -------------------------------------------------------------------------------- /src/partition.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use crate::error::*; 4 | use crate::{ 5 | client::Client, 6 | proto::{ 7 | self, 8 | common::{MsgBase, MsgType}, 9 | }, 10 | utils::status_to_result, 11 | }; 12 | 13 | impl Client { 14 | pub async fn create_partition( 15 | &self, 16 | collection_name: String, 17 | partition_name: String, 18 | ) -> Result<()> { 19 | status_to_result(&Some( 20 | self.client 21 | .clone() 22 | .create_partition(crate::proto::milvus::CreatePartitionRequest { 23 | base: Some(MsgBase::new(MsgType::CreatePartition)), 24 | db_name: "".to_string(), // reserved 25 | collection_name, 26 | partition_name, 27 | }) 28 | .await? 29 | .into_inner(), 30 | )) 31 | } 32 | 33 | pub async fn drop_partition( 34 | &self, 35 | collection_name: String, 36 | partition_name: String, 37 | ) -> Result<()> { 38 | status_to_result(&Some( 39 | self.client 40 | .clone() 41 | .drop_partition(crate::proto::milvus::DropPartitionRequest { 42 | base: Some(MsgBase::new(MsgType::DropPartition)), 43 | db_name: "".to_string(), // reserved 44 | collection_name, 45 | partition_name, 46 | }) 47 | .await? 48 | .into_inner(), 49 | )) 50 | } 51 | 52 | pub async fn list_partitions(&self, collection_name: String) -> Result> { 53 | let res = self 54 | .client 55 | .clone() 56 | .show_partitions(crate::proto::milvus::ShowPartitionsRequest { 57 | base: Some(MsgBase::new(MsgType::ShowPartitions)), 58 | db_name: "".to_string(), // reserved 59 | collection_name, 60 | collection_id: 0, // reserved 61 | partition_names: vec![], // reserved 62 | r#type: 0, // reserved 63 | }) 64 | .await? 65 | .into_inner(); 66 | status_to_result(&res.status)?; 67 | Ok(res.partition_names) 68 | } 69 | 70 | pub async fn has_partition( 71 | &self, 72 | collection_name: String, 73 | partition_name: String, 74 | ) -> Result { 75 | let res = self 76 | .client 77 | .clone() 78 | .has_partition(crate::proto::milvus::HasPartitionRequest { 79 | base: Some(MsgBase::new(MsgType::HasPartition)), 80 | db_name: "".to_string(), // reserved 81 | collection_name, 82 | partition_name, 83 | }) 84 | .await? 85 | .into_inner(); 86 | status_to_result(&res.status)?; 87 | Ok(res.value) 88 | } 89 | 90 | pub async fn get_partition_stats( 91 | &self, 92 | collection_name: String, 93 | partition_name: String, 94 | ) -> Result> { 95 | let res = self 96 | .client 97 | .clone() 98 | .get_partition_statistics(crate::proto::milvus::GetPartitionStatisticsRequest { 99 | base: Some(MsgBase::new(MsgType::GetPartitionStatistics)), 100 | db_name: "".to_string(), // reserved 101 | collection_name, 102 | partition_name, 103 | }) 104 | .await? 105 | .into_inner(); 106 | status_to_result(&res.status)?; 107 | 108 | Ok(res.stats.into_iter().map(|s| (s.key, s.value)).collect()) 109 | } 110 | 111 | // pub async fn load_partitions, I: IntoIterator>( 112 | // &self, 113 | // collection_name: S, 114 | // partition_names: I, 115 | // replica_number: i32, 116 | // ) -> Result<()> { 117 | // let names: Vec = partition_names.into_iter().map(|x| x.to_string()).collect(); 118 | // status_to_result(&Some( 119 | // self.client 120 | // .clone() 121 | // .load_partitions(proto::milvus::LoadPartitionsRequest { 122 | // base: Some(MsgBase::new(MsgType::LoadPartitions)), 123 | // db_name: "".to_string(), 124 | // collection_name: collection_name.into(), 125 | // replica_number, 126 | // partition_names: names.clone(), 127 | // }) 128 | // .await? 129 | // .into_inner(), 130 | // ))?; 131 | 132 | // loop { 133 | // if self.get_loading_progress(&names).await? >= 100 { 134 | // return Ok(()); 135 | // } 136 | 137 | // tokio::time::sleep(Duration::from_millis(config::WAIT_LOAD_DURATION_MS)).await; 138 | // } 139 | // } 140 | 141 | // pub async fn release_partitions>( 142 | // &self, 143 | // partition_names: I, 144 | // ) -> Result<()> { 145 | // status_to_result(&Some( 146 | // self.client 147 | // .clone() 148 | // .release_partitions(ReleasePartitionsRequest { 149 | // base: Some(MsgBase::new(MsgType::ReleasePartitions)), 150 | // db_name: "".to_string(), 151 | // collection_name: self.schema().name.to_string(), 152 | // partition_names: partition_names.into_iter().map(|x| x.to_string()).collect(), 153 | // }) 154 | // .await? 155 | // .into_inner(), 156 | // )) 157 | // } 158 | } 159 | -------------------------------------------------------------------------------- /src/proto/milvus.proto.feder.rs: -------------------------------------------------------------------------------- 1 | #[allow(clippy::derive_partial_eq_without_eq)] 2 | #[derive(Clone, PartialEq, ::prost::Message)] 3 | pub struct SegmentIndexData { 4 | #[prost(int64, tag = "1")] 5 | pub segment_id: i64, 6 | /// data from knownwhere 7 | #[prost(string, tag = "2")] 8 | pub index_data: ::prost::alloc::string::String, 9 | } 10 | #[allow(clippy::derive_partial_eq_without_eq)] 11 | #[derive(Clone, PartialEq, ::prost::Message)] 12 | pub struct FederSegmentSearchResult { 13 | #[prost(int64, tag = "1")] 14 | pub segment_id: i64, 15 | #[prost(string, tag = "2")] 16 | pub visit_info: ::prost::alloc::string::String, 17 | } 18 | #[allow(clippy::derive_partial_eq_without_eq)] 19 | #[derive(Clone, PartialEq, ::prost::Message)] 20 | pub struct ListIndexedSegmentRequest { 21 | #[prost(message, optional, tag = "1")] 22 | pub base: ::core::option::Option, 23 | #[prost(string, tag = "2")] 24 | pub collection_name: ::prost::alloc::string::String, 25 | #[prost(string, tag = "3")] 26 | pub index_name: ::prost::alloc::string::String, 27 | } 28 | #[allow(clippy::derive_partial_eq_without_eq)] 29 | #[derive(Clone, PartialEq, ::prost::Message)] 30 | pub struct ListIndexedSegmentResponse { 31 | #[prost(message, optional, tag = "1")] 32 | pub status: ::core::option::Option, 33 | #[prost(int64, repeated, tag = "2")] 34 | pub segment_i_ds: ::prost::alloc::vec::Vec, 35 | } 36 | #[allow(clippy::derive_partial_eq_without_eq)] 37 | #[derive(Clone, PartialEq, ::prost::Message)] 38 | pub struct DescribeSegmentIndexDataRequest { 39 | #[prost(message, optional, tag = "1")] 40 | pub base: ::core::option::Option, 41 | #[prost(string, tag = "2")] 42 | pub collection_name: ::prost::alloc::string::String, 43 | #[prost(string, tag = "3")] 44 | pub index_name: ::prost::alloc::string::String, 45 | #[prost(int64, repeated, tag = "4")] 46 | pub segments_i_ds: ::prost::alloc::vec::Vec, 47 | } 48 | #[allow(clippy::derive_partial_eq_without_eq)] 49 | #[derive(Clone, PartialEq, ::prost::Message)] 50 | pub struct DescribeSegmentIndexDataResponse { 51 | #[prost(message, optional, tag = "1")] 52 | pub status: ::core::option::Option, 53 | /// segmentID => segmentIndexData 54 | #[prost(map = "int64, message", tag = "2")] 55 | pub index_data: ::std::collections::HashMap, 56 | #[prost(message, repeated, tag = "3")] 57 | pub index_params: ::prost::alloc::vec::Vec, 58 | } 59 | -------------------------------------------------------------------------------- /src/proto/milvus.proto.msg.rs: -------------------------------------------------------------------------------- 1 | #[allow(clippy::derive_partial_eq_without_eq)] 2 | #[derive(Clone, PartialEq, ::prost::Message)] 3 | pub struct InsertRequest { 4 | #[prost(message, optional, tag = "1")] 5 | pub base: ::core::option::Option, 6 | #[prost(string, tag = "2")] 7 | pub shard_name: ::prost::alloc::string::String, 8 | #[prost(string, tag = "3")] 9 | pub db_name: ::prost::alloc::string::String, 10 | #[prost(string, tag = "4")] 11 | pub collection_name: ::prost::alloc::string::String, 12 | #[prost(string, tag = "5")] 13 | pub partition_name: ::prost::alloc::string::String, 14 | #[prost(int64, tag = "6")] 15 | pub db_id: i64, 16 | #[prost(int64, tag = "7")] 17 | pub collection_id: i64, 18 | #[prost(int64, tag = "8")] 19 | pub partition_id: i64, 20 | #[prost(int64, tag = "9")] 21 | pub segment_id: i64, 22 | #[prost(uint64, repeated, tag = "10")] 23 | pub timestamps: ::prost::alloc::vec::Vec, 24 | #[prost(int64, repeated, tag = "11")] 25 | pub row_i_ds: ::prost::alloc::vec::Vec, 26 | /// row_data was reserved for compatibility 27 | #[prost(message, repeated, tag = "12")] 28 | pub row_data: ::prost::alloc::vec::Vec, 29 | #[prost(message, repeated, tag = "13")] 30 | pub fields_data: ::prost::alloc::vec::Vec, 31 | #[prost(uint64, tag = "14")] 32 | pub num_rows: u64, 33 | #[prost(enumeration = "InsertDataVersion", tag = "15")] 34 | pub version: i32, 35 | } 36 | #[allow(clippy::derive_partial_eq_without_eq)] 37 | #[derive(Clone, PartialEq, ::prost::Message)] 38 | pub struct DeleteRequest { 39 | #[prost(message, optional, tag = "1")] 40 | pub base: ::core::option::Option, 41 | #[prost(string, tag = "2")] 42 | pub shard_name: ::prost::alloc::string::String, 43 | #[prost(string, tag = "3")] 44 | pub db_name: ::prost::alloc::string::String, 45 | #[prost(string, tag = "4")] 46 | pub collection_name: ::prost::alloc::string::String, 47 | #[prost(string, tag = "5")] 48 | pub partition_name: ::prost::alloc::string::String, 49 | #[prost(int64, tag = "6")] 50 | pub db_id: i64, 51 | #[prost(int64, tag = "7")] 52 | pub collection_id: i64, 53 | #[prost(int64, tag = "8")] 54 | pub partition_id: i64, 55 | /// deprecated 56 | #[prost(int64, repeated, tag = "9")] 57 | pub int64_primary_keys: ::prost::alloc::vec::Vec, 58 | #[prost(uint64, repeated, tag = "10")] 59 | pub timestamps: ::prost::alloc::vec::Vec, 60 | #[prost(int64, tag = "11")] 61 | pub num_rows: i64, 62 | #[prost(message, optional, tag = "12")] 63 | pub primary_keys: ::core::option::Option, 64 | } 65 | #[allow(clippy::derive_partial_eq_without_eq)] 66 | #[derive(Clone, PartialEq, ::prost::Message)] 67 | pub struct MsgPosition { 68 | #[prost(string, tag = "1")] 69 | pub channel_name: ::prost::alloc::string::String, 70 | #[prost(bytes = "vec", tag = "2")] 71 | pub msg_id: ::prost::alloc::vec::Vec, 72 | #[prost(string, tag = "3")] 73 | pub msg_group: ::prost::alloc::string::String, 74 | #[prost(uint64, tag = "4")] 75 | pub timestamp: u64, 76 | } 77 | #[allow(clippy::derive_partial_eq_without_eq)] 78 | #[derive(Clone, PartialEq, ::prost::Message)] 79 | pub struct CreateCollectionRequest { 80 | #[prost(message, optional, tag = "1")] 81 | pub base: ::core::option::Option, 82 | #[prost(string, tag = "2")] 83 | pub db_name: ::prost::alloc::string::String, 84 | #[prost(string, tag = "3")] 85 | pub collection_name: ::prost::alloc::string::String, 86 | #[prost(string, tag = "4")] 87 | pub partition_name: ::prost::alloc::string::String, 88 | /// `schema` is the serialized `schema.CollectionSchema` 89 | #[prost(int64, tag = "5")] 90 | pub db_id: i64, 91 | #[prost(int64, tag = "6")] 92 | pub collection_id: i64, 93 | /// deprecated 94 | #[prost(int64, tag = "7")] 95 | pub partition_id: i64, 96 | #[prost(bytes = "vec", tag = "8")] 97 | pub schema: ::prost::alloc::vec::Vec, 98 | #[prost(string, repeated, tag = "9")] 99 | pub virtual_channel_names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, 100 | #[prost(string, repeated, tag = "10")] 101 | pub physical_channel_names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, 102 | #[prost(int64, repeated, tag = "11")] 103 | pub partition_i_ds: ::prost::alloc::vec::Vec, 104 | } 105 | #[allow(clippy::derive_partial_eq_without_eq)] 106 | #[derive(Clone, PartialEq, ::prost::Message)] 107 | pub struct DropCollectionRequest { 108 | #[prost(message, optional, tag = "1")] 109 | pub base: ::core::option::Option, 110 | #[prost(string, tag = "2")] 111 | pub db_name: ::prost::alloc::string::String, 112 | #[prost(string, tag = "3")] 113 | pub collection_name: ::prost::alloc::string::String, 114 | #[prost(int64, tag = "4")] 115 | pub db_id: i64, 116 | #[prost(int64, tag = "5")] 117 | pub collection_id: i64, 118 | } 119 | #[allow(clippy::derive_partial_eq_without_eq)] 120 | #[derive(Clone, PartialEq, ::prost::Message)] 121 | pub struct CreatePartitionRequest { 122 | #[prost(message, optional, tag = "1")] 123 | pub base: ::core::option::Option, 124 | #[prost(string, tag = "2")] 125 | pub db_name: ::prost::alloc::string::String, 126 | #[prost(string, tag = "3")] 127 | pub collection_name: ::prost::alloc::string::String, 128 | #[prost(string, tag = "4")] 129 | pub partition_name: ::prost::alloc::string::String, 130 | #[prost(int64, tag = "5")] 131 | pub db_id: i64, 132 | #[prost(int64, tag = "6")] 133 | pub collection_id: i64, 134 | #[prost(int64, tag = "7")] 135 | pub partition_id: i64, 136 | } 137 | #[allow(clippy::derive_partial_eq_without_eq)] 138 | #[derive(Clone, PartialEq, ::prost::Message)] 139 | pub struct DropPartitionRequest { 140 | #[prost(message, optional, tag = "1")] 141 | pub base: ::core::option::Option, 142 | #[prost(string, tag = "2")] 143 | pub db_name: ::prost::alloc::string::String, 144 | #[prost(string, tag = "3")] 145 | pub collection_name: ::prost::alloc::string::String, 146 | #[prost(string, tag = "4")] 147 | pub partition_name: ::prost::alloc::string::String, 148 | #[prost(int64, tag = "5")] 149 | pub db_id: i64, 150 | #[prost(int64, tag = "6")] 151 | pub collection_id: i64, 152 | #[prost(int64, tag = "7")] 153 | pub partition_id: i64, 154 | } 155 | #[allow(clippy::derive_partial_eq_without_eq)] 156 | #[derive(Clone, PartialEq, ::prost::Message)] 157 | pub struct TimeTickMsg { 158 | #[prost(message, optional, tag = "1")] 159 | pub base: ::core::option::Option, 160 | } 161 | #[allow(clippy::derive_partial_eq_without_eq)] 162 | #[derive(Clone, PartialEq, ::prost::Message)] 163 | pub struct DataNodeTtMsg { 164 | #[prost(message, optional, tag = "1")] 165 | pub base: ::core::option::Option, 166 | #[prost(string, tag = "2")] 167 | pub channel_name: ::prost::alloc::string::String, 168 | #[prost(uint64, tag = "3")] 169 | pub timestamp: u64, 170 | #[prost(message, repeated, tag = "4")] 171 | pub segments_stats: ::prost::alloc::vec::Vec, 172 | } 173 | #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] 174 | #[repr(i32)] 175 | pub enum InsertDataVersion { 176 | /// 0 must refer to row-based format, since it's the first version in Milvus. 177 | RowBased = 0, 178 | ColumnBased = 1, 179 | } 180 | impl InsertDataVersion { 181 | /// String value of the enum field names used in the ProtoBuf definition. 182 | /// 183 | /// The values are not transformed in any way and thus are considered stable 184 | /// (if the ProtoBuf definition does not change) and safe for programmatic use. 185 | pub fn as_str_name(&self) -> &'static str { 186 | match self { 187 | InsertDataVersion::RowBased => "RowBased", 188 | InsertDataVersion::ColumnBased => "ColumnBased", 189 | } 190 | } 191 | /// Creates an enum from field names used in the ProtoBuf definition. 192 | pub fn from_str_name(value: &str) -> ::core::option::Option { 193 | match value { 194 | "RowBased" => Some(Self::RowBased), 195 | "ColumnBased" => Some(Self::ColumnBased), 196 | _ => None, 197 | } 198 | } 199 | } 200 | -------------------------------------------------------------------------------- /src/proto/milvus.proto.schema.rs: -------------------------------------------------------------------------------- 1 | /// * 2 | /// @brief Field schema 3 | #[allow(clippy::derive_partial_eq_without_eq)] 4 | #[derive(Clone, PartialEq, ::prost::Message)] 5 | pub struct FieldSchema { 6 | #[prost(int64, tag = "1")] 7 | pub field_id: i64, 8 | #[prost(string, tag = "2")] 9 | pub name: ::prost::alloc::string::String, 10 | #[prost(bool, tag = "3")] 11 | pub is_primary_key: bool, 12 | #[prost(string, tag = "4")] 13 | pub description: ::prost::alloc::string::String, 14 | #[prost(enumeration = "DataType", tag = "5")] 15 | pub data_type: i32, 16 | #[prost(message, repeated, tag = "6")] 17 | pub type_params: ::prost::alloc::vec::Vec, 18 | #[prost(message, repeated, tag = "7")] 19 | pub index_params: ::prost::alloc::vec::Vec, 20 | #[prost(bool, tag = "8")] 21 | pub auto_id: bool, 22 | /// To keep compatible with older version, the default 23 | #[prost(enumeration = "FieldState", tag = "9")] 24 | pub state: i32, 25 | /// state is `Created`. 26 | /// 27 | /// For array type, the element type is stored here 28 | #[prost(enumeration = "DataType", tag = "10")] 29 | pub element_type: i32, 30 | /// default_value only support scalars except array and json for now 31 | #[prost(message, optional, tag = "11")] 32 | pub default_value: ::core::option::Option, 33 | /// mark whether this field is the dynamic field 34 | #[prost(bool, tag = "12")] 35 | pub is_dynamic: bool, 36 | /// enable logic partitions 37 | #[prost(bool, tag = "13")] 38 | pub is_partition_key: bool, 39 | } 40 | /// * 41 | /// @brief Collection schema 42 | #[allow(clippy::derive_partial_eq_without_eq)] 43 | #[derive(Clone, PartialEq, ::prost::Message)] 44 | pub struct CollectionSchema { 45 | #[prost(string, tag = "1")] 46 | pub name: ::prost::alloc::string::String, 47 | #[prost(string, tag = "2")] 48 | pub description: ::prost::alloc::string::String, 49 | /// deprecated later, keep compatible with c++ part now 50 | #[deprecated] 51 | #[prost(bool, tag = "3")] 52 | pub auto_id: bool, 53 | #[prost(message, repeated, tag = "4")] 54 | pub fields: ::prost::alloc::vec::Vec, 55 | /// mark whether this table has the dynamic field function enabled. 56 | #[prost(bool, tag = "5")] 57 | pub enable_dynamic_field: bool, 58 | } 59 | #[allow(clippy::derive_partial_eq_without_eq)] 60 | #[derive(Clone, PartialEq, ::prost::Message)] 61 | pub struct BoolArray { 62 | #[prost(bool, repeated, tag = "1")] 63 | pub data: ::prost::alloc::vec::Vec, 64 | } 65 | #[allow(clippy::derive_partial_eq_without_eq)] 66 | #[derive(Clone, PartialEq, ::prost::Message)] 67 | pub struct IntArray { 68 | #[prost(int32, repeated, tag = "1")] 69 | pub data: ::prost::alloc::vec::Vec, 70 | } 71 | #[allow(clippy::derive_partial_eq_without_eq)] 72 | #[derive(Clone, PartialEq, ::prost::Message)] 73 | pub struct LongArray { 74 | #[prost(int64, repeated, tag = "1")] 75 | pub data: ::prost::alloc::vec::Vec, 76 | } 77 | #[allow(clippy::derive_partial_eq_without_eq)] 78 | #[derive(Clone, PartialEq, ::prost::Message)] 79 | pub struct FloatArray { 80 | #[prost(float, repeated, tag = "1")] 81 | pub data: ::prost::alloc::vec::Vec, 82 | } 83 | #[allow(clippy::derive_partial_eq_without_eq)] 84 | #[derive(Clone, PartialEq, ::prost::Message)] 85 | pub struct DoubleArray { 86 | #[prost(double, repeated, tag = "1")] 87 | pub data: ::prost::alloc::vec::Vec, 88 | } 89 | /// For special fields such as bigdecimal, array... 90 | #[allow(clippy::derive_partial_eq_without_eq)] 91 | #[derive(Clone, PartialEq, ::prost::Message)] 92 | pub struct BytesArray { 93 | #[prost(bytes = "vec", repeated, tag = "1")] 94 | pub data: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, 95 | } 96 | #[allow(clippy::derive_partial_eq_without_eq)] 97 | #[derive(Clone, PartialEq, ::prost::Message)] 98 | pub struct StringArray { 99 | #[prost(string, repeated, tag = "1")] 100 | pub data: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, 101 | } 102 | #[allow(clippy::derive_partial_eq_without_eq)] 103 | #[derive(Clone, PartialEq, ::prost::Message)] 104 | pub struct ArrayArray { 105 | #[prost(message, repeated, tag = "1")] 106 | pub data: ::prost::alloc::vec::Vec, 107 | #[prost(enumeration = "DataType", tag = "2")] 108 | pub element_type: i32, 109 | } 110 | #[allow(clippy::derive_partial_eq_without_eq)] 111 | #[derive(Clone, PartialEq, ::prost::Message)] 112 | pub struct JsonArray { 113 | #[prost(bytes = "vec", repeated, tag = "1")] 114 | pub data: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, 115 | } 116 | #[allow(clippy::derive_partial_eq_without_eq)] 117 | #[derive(Clone, PartialEq, ::prost::Message)] 118 | pub struct ValueField { 119 | #[prost(oneof = "value_field::Data", tags = "1, 2, 3, 4, 5, 6, 7")] 120 | pub data: ::core::option::Option, 121 | } 122 | /// Nested message and enum types in `ValueField`. 123 | pub mod value_field { 124 | #[allow(clippy::derive_partial_eq_without_eq)] 125 | #[derive(Clone, PartialEq, ::prost::Oneof)] 126 | pub enum Data { 127 | #[prost(bool, tag = "1")] 128 | BoolData(bool), 129 | #[prost(int32, tag = "2")] 130 | IntData(i32), 131 | #[prost(int64, tag = "3")] 132 | LongData(i64), 133 | #[prost(float, tag = "4")] 134 | FloatData(f32), 135 | #[prost(double, tag = "5")] 136 | DoubleData(f64), 137 | #[prost(string, tag = "6")] 138 | StringData(::prost::alloc::string::String), 139 | #[prost(bytes, tag = "7")] 140 | BytesData(::prost::alloc::vec::Vec), 141 | } 142 | } 143 | #[allow(clippy::derive_partial_eq_without_eq)] 144 | #[derive(Clone, PartialEq, ::prost::Message)] 145 | pub struct ScalarField { 146 | #[prost(oneof = "scalar_field::Data", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9")] 147 | pub data: ::core::option::Option, 148 | } 149 | /// Nested message and enum types in `ScalarField`. 150 | pub mod scalar_field { 151 | #[allow(clippy::derive_partial_eq_without_eq)] 152 | #[derive(Clone, PartialEq, ::prost::Oneof)] 153 | pub enum Data { 154 | #[prost(message, tag = "1")] 155 | BoolData(super::BoolArray), 156 | #[prost(message, tag = "2")] 157 | IntData(super::IntArray), 158 | #[prost(message, tag = "3")] 159 | LongData(super::LongArray), 160 | #[prost(message, tag = "4")] 161 | FloatData(super::FloatArray), 162 | #[prost(message, tag = "5")] 163 | DoubleData(super::DoubleArray), 164 | #[prost(message, tag = "6")] 165 | StringData(super::StringArray), 166 | #[prost(message, tag = "7")] 167 | BytesData(super::BytesArray), 168 | #[prost(message, tag = "8")] 169 | ArrayData(super::ArrayArray), 170 | #[prost(message, tag = "9")] 171 | JsonData(super::JsonArray), 172 | } 173 | } 174 | #[allow(clippy::derive_partial_eq_without_eq)] 175 | #[derive(Clone, PartialEq, ::prost::Message)] 176 | pub struct VectorField { 177 | #[prost(int64, tag = "1")] 178 | pub dim: i64, 179 | #[prost(oneof = "vector_field::Data", tags = "2, 3, 4, 5")] 180 | pub data: ::core::option::Option, 181 | } 182 | /// Nested message and enum types in `VectorField`. 183 | pub mod vector_field { 184 | #[allow(clippy::derive_partial_eq_without_eq)] 185 | #[derive(Clone, PartialEq, ::prost::Oneof)] 186 | pub enum Data { 187 | #[prost(message, tag = "2")] 188 | FloatVector(super::FloatArray), 189 | #[prost(bytes, tag = "3")] 190 | BinaryVector(::prost::alloc::vec::Vec), 191 | #[prost(bytes, tag = "4")] 192 | Float16Vector(::prost::alloc::vec::Vec), 193 | #[prost(bytes, tag = "5")] 194 | Bfloat16Vector(::prost::alloc::vec::Vec), 195 | } 196 | } 197 | #[allow(clippy::derive_partial_eq_without_eq)] 198 | #[derive(Clone, PartialEq, ::prost::Message)] 199 | pub struct FieldData { 200 | #[prost(enumeration = "DataType", tag = "1")] 201 | pub r#type: i32, 202 | #[prost(string, tag = "2")] 203 | pub field_name: ::prost::alloc::string::String, 204 | #[prost(int64, tag = "5")] 205 | pub field_id: i64, 206 | #[prost(bool, tag = "6")] 207 | pub is_dynamic: bool, 208 | #[prost(oneof = "field_data::Field", tags = "3, 4")] 209 | pub field: ::core::option::Option, 210 | } 211 | /// Nested message and enum types in `FieldData`. 212 | pub mod field_data { 213 | #[allow(clippy::derive_partial_eq_without_eq)] 214 | #[derive(Clone, PartialEq, ::prost::Oneof)] 215 | pub enum Field { 216 | #[prost(message, tag = "3")] 217 | Scalars(super::ScalarField), 218 | #[prost(message, tag = "4")] 219 | Vectors(super::VectorField), 220 | } 221 | } 222 | #[allow(clippy::derive_partial_eq_without_eq)] 223 | #[derive(Clone, PartialEq, ::prost::Message)] 224 | pub struct IDs { 225 | #[prost(oneof = "i_ds::IdField", tags = "1, 2")] 226 | pub id_field: ::core::option::Option, 227 | } 228 | /// Nested message and enum types in `IDs`. 229 | pub mod i_ds { 230 | #[allow(clippy::derive_partial_eq_without_eq)] 231 | #[derive(Clone, PartialEq, ::prost::Oneof)] 232 | pub enum IdField { 233 | #[prost(message, tag = "1")] 234 | IntId(super::LongArray), 235 | #[prost(message, tag = "2")] 236 | StrId(super::StringArray), 237 | } 238 | } 239 | #[allow(clippy::derive_partial_eq_without_eq)] 240 | #[derive(Clone, PartialEq, ::prost::Message)] 241 | pub struct SearchResultData { 242 | #[prost(int64, tag = "1")] 243 | pub num_queries: i64, 244 | #[prost(int64, tag = "2")] 245 | pub top_k: i64, 246 | #[prost(message, repeated, tag = "3")] 247 | pub fields_data: ::prost::alloc::vec::Vec, 248 | #[prost(float, repeated, tag = "4")] 249 | pub scores: ::prost::alloc::vec::Vec, 250 | #[prost(message, optional, tag = "5")] 251 | pub ids: ::core::option::Option, 252 | #[prost(int64, repeated, tag = "6")] 253 | pub topks: ::prost::alloc::vec::Vec, 254 | #[prost(string, repeated, tag = "7")] 255 | pub output_fields: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, 256 | #[prost(message, optional, tag = "8")] 257 | pub group_by_field_value: ::core::option::Option, 258 | } 259 | /// * 260 | /// @brief Field data type 261 | #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] 262 | #[repr(i32)] 263 | pub enum DataType { 264 | None = 0, 265 | Bool = 1, 266 | Int8 = 2, 267 | Int16 = 3, 268 | Int32 = 4, 269 | Int64 = 5, 270 | Float = 10, 271 | Double = 11, 272 | String = 20, 273 | /// variable-length strings with a specified maximum length 274 | VarChar = 21, 275 | Array = 22, 276 | Json = 23, 277 | BinaryVector = 100, 278 | FloatVector = 101, 279 | Float16Vector = 102, 280 | BFloat16Vector = 103, 281 | } 282 | impl DataType { 283 | /// String value of the enum field names used in the ProtoBuf definition. 284 | /// 285 | /// The values are not transformed in any way and thus are considered stable 286 | /// (if the ProtoBuf definition does not change) and safe for programmatic use. 287 | pub fn as_str_name(&self) -> &'static str { 288 | match self { 289 | DataType::None => "None", 290 | DataType::Bool => "Bool", 291 | DataType::Int8 => "Int8", 292 | DataType::Int16 => "Int16", 293 | DataType::Int32 => "Int32", 294 | DataType::Int64 => "Int64", 295 | DataType::Float => "Float", 296 | DataType::Double => "Double", 297 | DataType::String => "String", 298 | DataType::VarChar => "VarChar", 299 | DataType::Array => "Array", 300 | DataType::Json => "JSON", 301 | DataType::BinaryVector => "BinaryVector", 302 | DataType::FloatVector => "FloatVector", 303 | DataType::Float16Vector => "Float16Vector", 304 | DataType::BFloat16Vector => "BFloat16Vector", 305 | } 306 | } 307 | /// Creates an enum from field names used in the ProtoBuf definition. 308 | pub fn from_str_name(value: &str) -> ::core::option::Option { 309 | match value { 310 | "None" => Some(Self::None), 311 | "Bool" => Some(Self::Bool), 312 | "Int8" => Some(Self::Int8), 313 | "Int16" => Some(Self::Int16), 314 | "Int32" => Some(Self::Int32), 315 | "Int64" => Some(Self::Int64), 316 | "Float" => Some(Self::Float), 317 | "Double" => Some(Self::Double), 318 | "String" => Some(Self::String), 319 | "VarChar" => Some(Self::VarChar), 320 | "Array" => Some(Self::Array), 321 | "JSON" => Some(Self::Json), 322 | "BinaryVector" => Some(Self::BinaryVector), 323 | "FloatVector" => Some(Self::FloatVector), 324 | "Float16Vector" => Some(Self::Float16Vector), 325 | "BFloat16Vector" => Some(Self::BFloat16Vector), 326 | _ => None, 327 | } 328 | } 329 | } 330 | #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] 331 | #[repr(i32)] 332 | pub enum FieldState { 333 | FieldCreated = 0, 334 | FieldCreating = 1, 335 | FieldDropping = 2, 336 | FieldDropped = 3, 337 | } 338 | impl FieldState { 339 | /// String value of the enum field names used in the ProtoBuf definition. 340 | /// 341 | /// The values are not transformed in any way and thus are considered stable 342 | /// (if the ProtoBuf definition does not change) and safe for programmatic use. 343 | pub fn as_str_name(&self) -> &'static str { 344 | match self { 345 | FieldState::FieldCreated => "FieldCreated", 346 | FieldState::FieldCreating => "FieldCreating", 347 | FieldState::FieldDropping => "FieldDropping", 348 | FieldState::FieldDropped => "FieldDropped", 349 | } 350 | } 351 | /// Creates an enum from field names used in the ProtoBuf definition. 352 | pub fn from_str_name(value: &str) -> ::core::option::Option { 353 | match value { 354 | "FieldCreated" => Some(Self::FieldCreated), 355 | "FieldCreating" => Some(Self::FieldCreating), 356 | "FieldDropping" => Some(Self::FieldDropping), 357 | "FieldDropped" => Some(Self::FieldDropped), 358 | _ => None, 359 | } 360 | } 361 | } 362 | -------------------------------------------------------------------------------- /src/proto/mod.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use self::common::{MsgBase, MsgType}; 18 | 19 | #[path = "milvus.proto.common.rs"] 20 | pub mod common; 21 | #[path = "milvus.proto.feder.rs"] 22 | pub mod feder; 23 | #[path = "milvus.proto.milvus.rs"] 24 | pub mod milvus; 25 | #[path = "milvus.proto.msg.rs"] 26 | pub mod msg; 27 | #[path = "milvus.proto.schema.rs"] 28 | pub mod schema; 29 | 30 | impl MsgBase { 31 | pub fn new(msg_type: MsgType) -> Self { 32 | Self { 33 | msg_type: msg_type.into(), 34 | ..Default::default() 35 | } 36 | } 37 | } 38 | -------------------------------------------------------------------------------- /src/query.rs: -------------------------------------------------------------------------------- 1 | use std::borrow::Borrow; 2 | use std::collections::HashMap; 3 | 4 | use prost::bytes::BytesMut; 5 | use prost::Message; 6 | 7 | use crate::client::{Client, ConsistencyLevel}; 8 | use crate::collection::{ParamValue, SearchResult}; 9 | use crate::data::FieldColumn; 10 | use crate::error::Error as SuperError; 11 | use crate::index::MetricType; 12 | use crate::proto::common::{ 13 | DslType, KeyValuePair, MsgBase, MsgType, PlaceholderGroup, PlaceholderType, PlaceholderValue, 14 | }; 15 | use crate::proto::milvus::SearchRequest; 16 | use crate::proto::schema::DataType; 17 | use crate::utils::status_to_result; 18 | use crate::value::Value; 19 | use crate::{error::*, proto}; 20 | 21 | const STRONG_TIMESTAMP: u64 = 0; 22 | const BOUNDED_TIMESTAMP: u64 = 2; 23 | const EVENTUALLY_TIMESTAMP: u64 = 1; 24 | 25 | #[derive(Debug, Clone)] 26 | pub struct QueryOptions { 27 | output_fields: Vec, 28 | partition_names: Vec, 29 | } 30 | 31 | impl Default for QueryOptions { 32 | fn default() -> Self { 33 | Self { 34 | output_fields: Vec::new(), 35 | partition_names: Vec::new(), 36 | } 37 | } 38 | } 39 | 40 | impl QueryOptions { 41 | pub fn new() -> Self { 42 | Self::default() 43 | } 44 | 45 | pub fn with_output_fields(output_fields: Vec) -> Self { 46 | Self::default().output_fields(output_fields) 47 | } 48 | 49 | pub fn with_partition_names(partition_names: Vec) -> Self { 50 | Self::default().partition_names(partition_names) 51 | } 52 | 53 | pub fn output_fields(mut self, output_fields: Vec) -> Self { 54 | self.output_fields = output_fields; 55 | self 56 | } 57 | 58 | pub fn partition_names(mut self, partition_names: Vec) -> Self { 59 | self.partition_names = partition_names; 60 | self 61 | } 62 | } 63 | 64 | pub struct SearchOptions { 65 | pub(crate) expr: String, 66 | pub(crate) limit: usize, 67 | pub(crate) output_fields: Vec, 68 | pub(crate) partitions: Vec, 69 | pub(crate) params: serde_json::Value, 70 | pub(crate) metric_type: MetricType, 71 | } 72 | 73 | impl Default for SearchOptions { 74 | fn default() -> Self { 75 | Self { 76 | expr: String::new(), 77 | limit: 10, 78 | output_fields: Vec::new(), 79 | partitions: Vec::new(), 80 | params: serde_json::Value::default(), 81 | metric_type: MetricType::L2, 82 | } 83 | } 84 | } 85 | 86 | impl SearchOptions { 87 | pub fn new() -> Self { 88 | Self::default() 89 | } 90 | 91 | pub fn with_expr(expr: String) -> Self { 92 | Self::default().expr(expr) 93 | } 94 | 95 | pub fn with_limit(limit: usize) -> Self { 96 | Self::default().limit(limit) 97 | } 98 | 99 | pub fn with_output_fields(output_fields: Vec) -> Self { 100 | Self::default().output_fields(output_fields) 101 | } 102 | 103 | pub fn with_partitions(partitions: Vec) -> Self { 104 | Self::default().partitions(partitions) 105 | } 106 | 107 | pub fn with_params(params: HashMap) -> Self { 108 | let mut options = Self::default(); 109 | for (k, v) in params { 110 | options = options.add_param(k, v); 111 | } 112 | options 113 | } 114 | 115 | pub fn with_metric_type(metric_type: MetricType) -> Self { 116 | Self::default().metric_type(metric_type) 117 | } 118 | 119 | pub fn radius(self, radius: f32) -> Self { 120 | self.add_param("radius", ParamValue!(radius)) 121 | } 122 | 123 | pub fn range_filter(self, filter: f32) -> Self { 124 | self.add_param("range_filter", ParamValue!(filter)) 125 | } 126 | 127 | pub fn expr(mut self, expr: String) -> Self { 128 | self.expr = expr; 129 | self 130 | } 131 | 132 | pub fn limit(mut self, limit: usize) -> Self { 133 | self.limit = limit; 134 | self 135 | } 136 | 137 | pub fn output_fields(mut self, output_fields: Vec) -> Self { 138 | self.output_fields = output_fields; 139 | self 140 | } 141 | 142 | pub fn partitions(mut self, partitions: Vec) -> Self { 143 | self.partitions = partitions; 144 | self 145 | } 146 | 147 | pub fn add_param(mut self, key: impl Into, value: ParamValue) -> Self { 148 | if self.params.is_null() { 149 | self.params = ParamValue!({}); 150 | } 151 | self.params 152 | .as_object_mut() 153 | .unwrap() 154 | .insert(key.into(), value); 155 | self 156 | } 157 | 158 | pub fn metric_type(mut self, metric_type: MetricType) -> Self { 159 | self.metric_type = metric_type; 160 | self 161 | } 162 | } 163 | 164 | impl Client { 165 | pub(crate) async fn get_gts_from_consistency( 166 | &self, 167 | collection_name: &str, 168 | consistency_level: ConsistencyLevel, 169 | ) -> u64 { 170 | match consistency_level { 171 | ConsistencyLevel::Strong => STRONG_TIMESTAMP, 172 | ConsistencyLevel::Bounded => BOUNDED_TIMESTAMP, 173 | ConsistencyLevel::Eventually => EVENTUALLY_TIMESTAMP, 174 | ConsistencyLevel::Session => self 175 | .collection_cache 176 | .get_timestamp(collection_name) 177 | .unwrap_or(EVENTUALLY_TIMESTAMP), 178 | 179 | // This level not works for now 180 | ConsistencyLevel::Customized => 0, 181 | } 182 | } 183 | 184 | pub async fn query( 185 | &self, 186 | collection_name: S, 187 | expr: Exp, 188 | options: &QueryOptions, 189 | ) -> Result> 190 | where 191 | S: AsRef, 192 | Exp: AsRef, 193 | { 194 | let collection_name = collection_name.as_ref(); 195 | let collection = self.collection_cache.get(collection_name).await?; 196 | 197 | let consistency_level = collection.consistency_level; 198 | let mut output_fields = options.output_fields.clone(); 199 | if output_fields.is_empty() { 200 | output_fields = collection.fields.iter().map(|f| f.name.clone()).collect(); 201 | } 202 | 203 | let res = self 204 | .client 205 | .clone() 206 | .query(proto::milvus::QueryRequest { 207 | base: Some(MsgBase::new(MsgType::Retrieve)), 208 | db_name: "".to_owned(), 209 | collection_name: collection_name.to_owned(), 210 | expr: expr.as_ref().to_owned(), 211 | output_fields: output_fields, 212 | partition_names: options.partition_names.clone(), 213 | travel_timestamp: 0, 214 | guarantee_timestamp: self 215 | .get_gts_from_consistency(collection_name, consistency_level) 216 | .await, 217 | query_params: Vec::new(), 218 | not_return_all_meta: false, 219 | consistency_level: ConsistencyLevel::default() as _, 220 | use_default_consistency: false, 221 | }) 222 | .await? 223 | .into_inner(); 224 | 225 | status_to_result(&res.status)?; 226 | 227 | Ok(res 228 | .fields_data 229 | .into_iter() 230 | .map(|f| FieldColumn::from(f)) 231 | .collect()) 232 | } 233 | 234 | pub async fn search( 235 | &self, 236 | collection_name: S, 237 | data: Vec>, 238 | vec_field: S, 239 | option: &SearchOptions, 240 | ) -> Result>> 241 | where 242 | S: Into, 243 | { 244 | // check and prepare params 245 | let search_params: Vec = vec![ 246 | KeyValuePair { 247 | key: "anns_field".to_owned(), 248 | value: vec_field.into(), 249 | }, 250 | KeyValuePair { 251 | key: "topk".to_owned(), 252 | value: option.limit.to_string(), 253 | }, 254 | KeyValuePair { 255 | key: "params".to_owned(), 256 | value: serde_json::to_string(&option.params)?, 257 | }, 258 | KeyValuePair { 259 | key: "metric_type".to_owned(), 260 | value: option.metric_type.to_string(), 261 | }, 262 | KeyValuePair { 263 | key: "round_decimal".to_owned(), 264 | value: "-1".to_owned(), 265 | }, 266 | ]; 267 | 268 | let collection_name = collection_name.into(); 269 | let collection = self.collection_cache.get(&collection_name).await?; 270 | 271 | let res = self 272 | .client 273 | .clone() 274 | .search(SearchRequest { 275 | base: Some(MsgBase::new(MsgType::Search)), 276 | db_name: "".to_string(), 277 | collection_name: collection_name.clone(), 278 | partition_names: option.partitions.clone(), 279 | dsl: option.expr.clone(), 280 | nq: data.len() as _, 281 | placeholder_group: get_place_holder_group(data)?, 282 | dsl_type: DslType::BoolExprV1 as _, 283 | output_fields: option 284 | .output_fields 285 | .clone() 286 | .into_iter() 287 | .map(|f| f.into()) 288 | .collect(), 289 | search_params, 290 | travel_timestamp: 0, 291 | guarantee_timestamp: self 292 | .get_gts_from_consistency(&collection_name, collection.consistency_level) 293 | .await, 294 | not_return_all_meta: false, 295 | consistency_level: ConsistencyLevel::default() as _, 296 | use_default_consistency: false, 297 | search_by_primary_keys: false, 298 | }) 299 | .await? 300 | .into_inner(); 301 | status_to_result(&res.status)?; 302 | let raw_data = res 303 | .results 304 | .ok_or(SuperError::Unexpected("no result for search".to_owned()))?; 305 | let mut result = Vec::new(); 306 | let mut offset = 0; 307 | let fields_data = raw_data 308 | .fields_data 309 | .into_iter() 310 | .map(Into::into) 311 | .collect::>(); 312 | let raw_id = raw_data.ids.unwrap().id_field.unwrap(); 313 | 314 | for k in raw_data.topks { 315 | let k = k as usize; 316 | let mut score = Vec::new(); 317 | score.extend_from_slice(&raw_data.scores[offset..offset + k]); 318 | let mut result_data = fields_data 319 | .iter() 320 | .map(FieldColumn::copy_with_metadata) 321 | .collect::>(); 322 | for j in 0..fields_data.len() { 323 | for i in offset..offset + k { 324 | result_data[j].push(fields_data[j].get(i).ok_or(SuperError::Unexpected( 325 | "out of range while indexing field data".to_owned(), 326 | ))?); 327 | } 328 | } 329 | 330 | let id = match raw_id { 331 | proto::schema::i_ds::IdField::IntId(ref d) => { 332 | Vec::::from_iter(d.data[offset..offset + k].iter().map(|&x| x.into())) 333 | } 334 | proto::schema::i_ds::IdField::StrId(ref d) => Vec::::from_iter( 335 | d.data[offset..offset + k].iter().map(|x| x.clone().into()), 336 | ), 337 | }; 338 | 339 | result.push(SearchResult { 340 | size: k as i64, 341 | score, 342 | field: result_data, 343 | id, 344 | }); 345 | 346 | offset += k; 347 | } 348 | 349 | Ok(result) 350 | } 351 | } 352 | 353 | fn get_place_holder_group(vectors: Vec) -> Result> { 354 | let group = PlaceholderGroup { 355 | placeholders: vec![get_place_holder_value(vectors)?], 356 | }; 357 | let mut buf = BytesMut::new(); 358 | group.encode(&mut buf).unwrap(); 359 | return Ok(buf.to_vec()); 360 | } 361 | 362 | fn get_place_holder_value(vectors: Vec) -> Result { 363 | let mut place_holder = PlaceholderValue { 364 | tag: "$0".to_string(), 365 | r#type: PlaceholderType::None as _, 366 | values: Vec::new(), 367 | }; 368 | // if no vectors, return an empty one 369 | if vectors.len() == 0 { 370 | return Ok(place_holder); 371 | }; 372 | 373 | match vectors[0] { 374 | Value::FloatArray(_) => place_holder.r#type = PlaceholderType::FloatVector as _, 375 | Value::Binary(_) => place_holder.r#type = PlaceholderType::BinaryVector as _, 376 | _ => { 377 | return Err(SuperError::from(crate::collection::Error::IllegalType( 378 | "place holder".to_string(), 379 | vec![DataType::BinaryVector, DataType::FloatVector], 380 | ))) 381 | } 382 | }; 383 | 384 | for v in &vectors { 385 | match (v, &vectors[0]) { 386 | (Value::FloatArray(d), Value::FloatArray(_)) => { 387 | let mut bytes = Vec::::with_capacity(d.len() * 4); 388 | for f in d.iter() { 389 | bytes.extend_from_slice(&f.to_le_bytes()); 390 | } 391 | place_holder.values.push(bytes) 392 | } 393 | (Value::Binary(d), Value::Binary(_)) => place_holder.values.push(d.to_vec()), 394 | _ => { 395 | return Err(SuperError::from(crate::collection::Error::IllegalType( 396 | "place holder".to_string(), 397 | vec![DataType::BinaryVector, DataType::FloatVector], 398 | ))) 399 | } 400 | }; 401 | } 402 | return Ok(place_holder); 403 | } 404 | -------------------------------------------------------------------------------- /src/schema.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use crate::error; 18 | use crate::error::Result; 19 | use crate::proto::schema::FieldState; 20 | use prost::alloc::vec::Vec; 21 | use prost::encoding::bool; 22 | use thiserror::Error as ThisError; 23 | 24 | use crate::proto::{ 25 | common::KeyValuePair, 26 | schema::{self, DataType}, 27 | }; 28 | 29 | pub use crate::proto::schema::FieldData; 30 | 31 | pub trait Schema { 32 | // fn name(&self) -> &str; 33 | // fn description(&self) -> &str; 34 | // fn fields(&self) -> &Vec; 35 | 36 | // fn schema(&self) -> CollectionSchema { 37 | // CollectionSchema { 38 | // name: self.name(), 39 | // description: self.description(), 40 | // fields: self.fields().to_owned(), 41 | // } 42 | // } 43 | 44 | // type ColumnIntoIter<'a>: Iterator)>; 45 | // type ColumnIter<'a>: Iterator, Value<'a>)>; 46 | 47 | // fn iter(&self) -> Self::ColumnIntoIter; // Self::ColumnIter<'_> 48 | // fn into_iter(self) -> Self::ColumnIntoIter; 49 | 50 | // fn validate(&self) -> std::result::Result<(), Error> { 51 | // for (schm, val) in self.iter() { 52 | // let dtype = val.data_type(); 53 | 54 | // if dtype != schm.dtype 55 | // && !(dtype == DataType::String && schm.dtype == DataType::VarChar) 56 | // { 57 | // return Err(Error::FieldWrongType( 58 | // schm.name.to_string(), 59 | // schm.dtype, 60 | // val.data_type(), 61 | // )); 62 | // } 63 | 64 | // match schm.dtype { 65 | // DataType::VarChar => match &val { 66 | // Value::String(d) if d.len() > schm.max_length as _ => { 67 | // return Err(Error::DimensionMismatch( 68 | // schm.name.to_string(), 69 | // schm.max_length as _, 70 | // d.len() as _, 71 | // )); 72 | // } 73 | // _ => unreachable!(), 74 | // }, 75 | // DataType::BinaryVector => match &val { 76 | // Value::Binary(d) => { 77 | // return Err(Error::DimensionMismatch( 78 | // schm.name.to_string(), 79 | // schm.dim as _, 80 | // d.len() as _, 81 | // )); 82 | // } 83 | // _ => unreachable!(), 84 | // }, 85 | // DataType::FloatVector => match &val { 86 | // Value::FloatArray(d) => { 87 | // return Err(Error::DimensionMismatch( 88 | // schm.name.to_string(), 89 | // schm.dim as _, 90 | // d.len() as _, 91 | // )); 92 | // } 93 | // _ => unreachable!(), 94 | // }, 95 | // _ => (), 96 | // } 97 | // } 98 | 99 | // Ok(()) 100 | // } 101 | } 102 | 103 | pub trait FromDataFields: Sized { 104 | fn from_data_fields(fileds: Vec) -> Option; 105 | } 106 | 107 | // pub trait Column<'a>: IntoFieldData + FromDataFields { 108 | // type Entity: Schema; 109 | // type IterRows: Iterator + 'a; 110 | // type IterColumns: Iterator> + 'a; 111 | 112 | // fn index(&self, idx: usize) -> Option; 113 | // fn with_capacity(cap: usize) -> Self; 114 | // fn add(&mut self, entity: Self::Entity); 115 | // fn len(&self) -> usize; 116 | // fn iter_columns(&'a self) -> Self::IterColumns; 117 | 118 | // fn iter_rows(&self) -> Box + '_> { 119 | // Box::new((0..self.len()).filter_map(|idx| self.index(idx))) 120 | // } 121 | 122 | // fn is_empty(&self) -> bool { 123 | // self.len() == 0 124 | // } 125 | 126 | // fn columns() -> &'static [FieldSchema<'static>] { 127 | // Self::Entity::SCHEMA 128 | // } 129 | // } 130 | 131 | // Bool = 1, 132 | // Int8 = 2, 133 | // Int16 = 3, 134 | // Int32 = 4, 135 | // Int64 = 5, 136 | // Float = 10, 137 | // Double = 11, 138 | // String = 20, 139 | // /// variable-length strings with a specified maximum length 140 | // VarChar = 21, 141 | // BinaryVector = 100, 142 | // FloatVector = 101, 143 | 144 | pub trait IntoFieldData { 145 | fn into_data_fields(self) -> Vec; 146 | } 147 | 148 | #[derive(Debug, Clone)] 149 | pub struct FieldSchema { 150 | pub name: String, 151 | pub description: String, 152 | pub dtype: DataType, 153 | pub is_primary: bool, 154 | pub auto_id: bool, 155 | pub chunk_size: usize, 156 | pub dim: i64, // only for BinaryVector and FloatVector 157 | pub max_length: i32, // only for VarChar 158 | } 159 | 160 | impl FieldSchema { 161 | pub const fn const_default() -> Self { 162 | Self { 163 | name: String::new(), 164 | description: String::new(), 165 | dtype: DataType::None, 166 | is_primary: false, 167 | auto_id: false, 168 | chunk_size: 0, 169 | dim: 0, 170 | max_length: 0, 171 | } 172 | } 173 | } 174 | 175 | impl Default for FieldSchema { 176 | fn default() -> Self { 177 | Self::const_default() 178 | } 179 | } 180 | 181 | impl From for FieldSchema { 182 | fn from(fld: schema::FieldSchema) -> Self { 183 | let dim: i64 = fld 184 | .type_params 185 | .iter() 186 | .find(|k| &k.key == "dim") 187 | .and_then(|x| x.value.parse().ok()) 188 | .unwrap_or(1); 189 | 190 | let dtype = DataType::from_i32(fld.data_type).unwrap(); 191 | 192 | FieldSchema { 193 | name: fld.name, 194 | description: fld.description, 195 | dtype, 196 | is_primary: fld.is_primary_key, 197 | auto_id: fld.auto_id, 198 | max_length: 0, 199 | chunk_size: (dim 200 | * match dtype { 201 | DataType::BinaryVector => dim / 8, 202 | _ => dim, 203 | }) as _, 204 | dim, 205 | } 206 | } 207 | } 208 | 209 | impl FieldSchema { 210 | pub fn new_bool(name: &str, description: &str) -> Self { 211 | Self { 212 | name: name.to_owned(), 213 | description: description.to_owned(), 214 | dtype: DataType::Bool, 215 | is_primary: false, 216 | auto_id: false, 217 | chunk_size: 1, 218 | dim: 1, 219 | max_length: 0, 220 | } 221 | } 222 | 223 | pub fn new_int8(name: &str, description: &str) -> Self { 224 | Self { 225 | name: name.to_owned(), 226 | description: description.to_owned(), 227 | dtype: DataType::Int8, 228 | is_primary: false, 229 | auto_id: false, 230 | chunk_size: 1, 231 | dim: 1, 232 | max_length: 0, 233 | } 234 | } 235 | 236 | pub fn new_int16(name: &str, description: &str) -> Self { 237 | Self { 238 | name: name.to_owned(), 239 | description: description.to_owned(), 240 | dtype: DataType::Int16, 241 | is_primary: false, 242 | auto_id: false, 243 | chunk_size: 1, 244 | dim: 1, 245 | max_length: 0, 246 | } 247 | } 248 | 249 | pub fn new_int32(name: &str, description: &str) -> Self { 250 | Self { 251 | name: name.to_owned(), 252 | description: description.to_owned(), 253 | dtype: DataType::Int32, 254 | is_primary: false, 255 | auto_id: false, 256 | chunk_size: 1, 257 | dim: 1, 258 | max_length: 0, 259 | } 260 | } 261 | 262 | pub fn new_int64(name: &str, description: &str) -> Self { 263 | Self { 264 | name: name.to_owned(), 265 | description: description.to_owned(), 266 | dtype: DataType::Int64, 267 | is_primary: false, 268 | auto_id: false, 269 | chunk_size: 1, 270 | dim: 1, 271 | max_length: 0, 272 | } 273 | } 274 | 275 | pub fn new_primary_int64(name: &str, description: &str, auto_id: bool) -> Self { 276 | Self { 277 | name: name.to_owned(), 278 | description: description.to_owned(), 279 | dtype: DataType::Int64, 280 | is_primary: true, 281 | auto_id, 282 | chunk_size: 1, 283 | dim: 1, 284 | max_length: 0, 285 | } 286 | } 287 | 288 | pub fn new_primary_varchar( 289 | name: &str, 290 | description: &str, 291 | auto_id: bool, 292 | max_length: i32, 293 | ) -> Self { 294 | Self { 295 | name: name.to_owned(), 296 | description: description.to_owned(), 297 | dtype: DataType::VarChar, 298 | is_primary: true, 299 | auto_id, 300 | max_length, 301 | chunk_size: 1, 302 | dim: 1, 303 | } 304 | } 305 | 306 | pub fn new_float(name: &str, description: &str) -> Self { 307 | Self { 308 | name: name.to_owned(), 309 | description: description.to_owned(), 310 | dtype: DataType::Float, 311 | is_primary: false, 312 | auto_id: false, 313 | chunk_size: 1, 314 | dim: 1, 315 | max_length: 0, 316 | } 317 | } 318 | 319 | pub fn new_double(name: &str, description: &str) -> Self { 320 | Self { 321 | name: name.to_owned(), 322 | description: description.to_owned(), 323 | dtype: DataType::Double, 324 | is_primary: false, 325 | auto_id: false, 326 | chunk_size: 1, 327 | dim: 1, 328 | max_length: 0, 329 | } 330 | } 331 | 332 | pub fn new_string(name: &str, description: &str) -> Self { 333 | Self { 334 | name: name.to_owned(), 335 | description: description.to_owned(), 336 | dtype: DataType::String, 337 | is_primary: false, 338 | auto_id: false, 339 | chunk_size: 1, 340 | dim: 1, 341 | max_length: 0, 342 | } 343 | } 344 | 345 | pub fn new_varchar(name: &str, description: &str, max_length: i32) -> Self { 346 | if max_length <= 0 { 347 | panic!("max_length should be positive"); 348 | } 349 | 350 | Self { 351 | name: name.to_owned(), 352 | description: description.to_owned(), 353 | dtype: DataType::VarChar, 354 | max_length, 355 | is_primary: false, 356 | auto_id: false, 357 | chunk_size: 1, 358 | dim: 1, 359 | } 360 | } 361 | 362 | pub fn new_binary_vector(name: &str, description: &str, dim: i64) -> Self { 363 | if dim <= 0 { 364 | panic!("dim should be positive"); 365 | } 366 | 367 | Self { 368 | name: name.to_owned(), 369 | description: description.to_owned(), 370 | dtype: DataType::BinaryVector, 371 | chunk_size: dim as usize / 8, 372 | dim, 373 | is_primary: false, 374 | auto_id: false, 375 | max_length: 0, 376 | } 377 | } 378 | 379 | pub fn new_float_vector(name: &str, description: &str, dim: i64) -> Self { 380 | if dim <= 0 { 381 | panic!("dim should be positive"); 382 | } 383 | 384 | Self { 385 | name: name.to_owned(), 386 | description: description.to_owned(), 387 | dtype: DataType::FloatVector, 388 | chunk_size: dim as usize, 389 | dim, 390 | is_primary: false, 391 | auto_id: false, 392 | max_length: 0, 393 | } 394 | } 395 | } 396 | 397 | impl From for schema::FieldSchema { 398 | fn from(fld: FieldSchema) -> schema::FieldSchema { 399 | let params = match fld.dtype { 400 | DataType::BinaryVector | DataType::FloatVector => vec![KeyValuePair { 401 | key: "dim".to_string(), 402 | value: fld.dim.to_string(), 403 | }], 404 | DataType::VarChar => vec![KeyValuePair { 405 | key: "max_length".to_string(), 406 | value: fld.max_length.to_string(), 407 | }], 408 | _ => Vec::new(), 409 | }; 410 | 411 | schema::FieldSchema { 412 | field_id: 0, 413 | name: fld.name.into(), 414 | is_primary_key: fld.is_primary, 415 | description: fld.description, 416 | data_type: fld.dtype as i32, 417 | type_params: params, 418 | index_params: Vec::new(), 419 | auto_id: fld.auto_id, 420 | state: FieldState::FieldCreated as _, 421 | element_type: 0, 422 | default_value: None, 423 | is_dynamic: false, 424 | is_partition_key: false, 425 | } 426 | } 427 | } 428 | 429 | #[derive(Debug, Clone)] 430 | pub struct CollectionSchema { 431 | pub(crate) name: String, 432 | pub(crate) description: String, 433 | pub(crate) fields: Vec, 434 | pub(crate) enable_dynamic_field: bool, 435 | } 436 | 437 | impl CollectionSchema { 438 | pub fn name(&self) -> &str { 439 | &self.name 440 | } 441 | 442 | #[inline] 443 | pub fn auto_id(&self) -> bool { 444 | self.fields.iter().any(|x| x.auto_id) 445 | } 446 | 447 | pub fn primary_column(&self) -> Option<&FieldSchema> { 448 | self.fields.iter().find(|s| s.is_primary) 449 | } 450 | 451 | pub fn validate(&self) -> Result<()> { 452 | self.primary_column().ok_or_else(|| Error::NoPrimaryKey)?; 453 | // TODO addidtional schema checks need to be added here 454 | Ok(()) 455 | } 456 | 457 | pub fn get_field(&self, name: S) -> Option<&FieldSchema> 458 | where 459 | S: AsRef, 460 | { 461 | let name = name.as_ref(); 462 | self.fields.iter().find(|f| f.name == name) 463 | } 464 | 465 | pub fn is_valid_vector_field(&self, field_name: &str) -> Result<()> { 466 | for f in &self.fields { 467 | if f.name == field_name { 468 | if f.dtype == DataType::BinaryVector || f.dtype == DataType::FloatVector { 469 | return Ok(()); 470 | } else { 471 | return Err(error::Error::from(Error::NotVectorField( 472 | field_name.to_owned(), 473 | ))); 474 | } 475 | } 476 | } 477 | return Err(error::Error::from(Error::NoSuchKey(field_name.to_owned()))); 478 | } 479 | } 480 | 481 | impl From for schema::CollectionSchema { 482 | fn from(col: CollectionSchema) -> Self { 483 | schema::CollectionSchema { 484 | name: col.name.to_string(), 485 | auto_id: col.auto_id(), 486 | description: col.description, 487 | fields: col.fields.into_iter().map(Into::into).collect(), 488 | enable_dynamic_field: col.enable_dynamic_field, 489 | } 490 | } 491 | } 492 | 493 | impl From for CollectionSchema { 494 | fn from(v: schema::CollectionSchema) -> Self { 495 | CollectionSchema { 496 | fields: v.fields.into_iter().map(Into::into).collect(), 497 | name: v.name, 498 | description: v.description, 499 | enable_dynamic_field: v.enable_dynamic_field, 500 | } 501 | } 502 | } 503 | 504 | #[derive(Debug, Clone)] 505 | pub struct CollectionSchemaBuilder { 506 | name: String, 507 | description: String, 508 | inner: Vec, 509 | enable_dynamic_field: bool, 510 | } 511 | 512 | impl CollectionSchemaBuilder { 513 | pub fn new(name: &str, description: &str) -> Self { 514 | Self { 515 | name: name.to_owned(), 516 | description: description.to_owned(), 517 | inner: Vec::new(), 518 | enable_dynamic_field: false, 519 | } 520 | } 521 | 522 | pub fn add_field(&mut self, schema: FieldSchema) -> &mut Self { 523 | self.inner.push(schema); 524 | self 525 | } 526 | 527 | pub fn set_primary_key(&mut self, name: S) -> Result<&mut Self> 528 | where 529 | S: AsRef, 530 | { 531 | let n = name.as_ref(); 532 | for f in self.inner.iter_mut() { 533 | if f.is_primary { 534 | return Err(error::Error::from(Error::DuplicatePrimaryKey( 535 | n.to_string(), 536 | f.name.to_string(), 537 | ))); 538 | } 539 | } 540 | 541 | for f in self.inner.iter_mut() { 542 | if n == f.name { 543 | if f.dtype == DataType::Int64 || f.dtype == DataType::VarChar { 544 | f.is_primary = true; 545 | return Ok(self); 546 | } else { 547 | return Err(error::Error::from(Error::UnsupportedPrimaryKey( 548 | f.dtype.to_owned(), 549 | ))); 550 | } 551 | } 552 | } 553 | 554 | Err(error::Error::from(Error::NoSuchKey(n.to_string()))) 555 | } 556 | 557 | pub fn enable_auto_id(&mut self) -> Result<&mut Self> { 558 | for f in self.inner.iter_mut() { 559 | if f.is_primary { 560 | if f.dtype == DataType::Int64 { 561 | f.auto_id = true; 562 | return Ok(self); 563 | } else { 564 | return Err(error::Error::from(Error::UnsupportedAutoId( 565 | f.dtype.to_owned(), 566 | ))); 567 | } 568 | } 569 | } 570 | 571 | Err(error::Error::from(Error::NoPrimaryKey)) 572 | } 573 | 574 | pub fn enable_dynamic_field(&mut self) -> &mut Self { 575 | self.enable_dynamic_field = true; 576 | self 577 | } 578 | 579 | pub fn build(&mut self) -> Result { 580 | let mut has_primary = false; 581 | 582 | for f in self.inner.iter() { 583 | if f.is_primary { 584 | has_primary = true; 585 | break; 586 | } 587 | } 588 | 589 | if !has_primary { 590 | return Err(error::Error::from(Error::NoPrimaryKey)); 591 | } 592 | 593 | let this = std::mem::replace(self, CollectionSchemaBuilder::new("".into(), "")); 594 | 595 | Ok(CollectionSchema { 596 | fields: this.inner.into(), 597 | name: this.name, 598 | description: this.description, 599 | enable_dynamic_field: self.enable_dynamic_field, 600 | }) 601 | } 602 | } 603 | 604 | #[derive(Debug, ThisError)] 605 | pub enum Error { 606 | #[error("try to set primary key {0:?}, but {1:?} is also key")] 607 | DuplicatePrimaryKey(String, String), 608 | 609 | #[error("can not find any primary key")] 610 | NoPrimaryKey, 611 | 612 | #[error("primary key must be int64 or varchar, unsupported type {0:?}")] 613 | UnsupportedPrimaryKey(DataType), 614 | 615 | #[error("auto id must be int64, unsupported type {0:?}")] 616 | UnsupportedAutoId(DataType), 617 | 618 | #[error("dimension mismatch for {0:?}, expected dim {1:?}, got {2:?}")] 619 | DimensionMismatch(String, i32, i32), 620 | 621 | #[error("wrong field data type, field {0} expected to be{1:?}, but got {2:?}")] 622 | FieldWrongType(String, DataType, DataType), 623 | 624 | #[error("field does not exists in schema: {0:?}")] 625 | FieldDoesNotExists(String), 626 | 627 | #[error("can not find such key {0:?}")] 628 | NoSuchKey(String), 629 | 630 | #[error("field {0:?} must be a vector field")] 631 | NotVectorField(String), 632 | } 633 | -------------------------------------------------------------------------------- /src/types.rs: -------------------------------------------------------------------------------- 1 | use crate::proto::{self, schema::DataType}; 2 | 3 | pub(crate) type Timestamp = u64; 4 | 5 | #[derive(Debug, Clone)] 6 | pub struct Field { 7 | pub id: i64, 8 | pub name: String, 9 | pub description: String, 10 | pub dtype: DataType, 11 | pub is_primary_key: bool, 12 | } 13 | 14 | impl From for Field { 15 | fn from(value: proto::schema::FieldSchema) -> Self { 16 | Self { 17 | id: value.field_id, 18 | name: value.name, 19 | description: value.description, 20 | dtype: DataType::from_i32(value.data_type).unwrap_or(DataType::None), 21 | is_primary_key: value.is_primary_key, 22 | } 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /src/utils.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use crate::{ 18 | error::Error, 19 | proto::common::{ErrorCode, Status}, 20 | }; 21 | 22 | pub fn status_to_result(status: &Option) -> Result<(), Error> { 23 | let status = status 24 | .clone() 25 | .ok_or(Error::Unexpected("no status".to_owned()))?; 26 | 27 | match ErrorCode::from_i32(status.error_code) { 28 | Some(i) => match i { 29 | ErrorCode::Success => Ok(()), 30 | _ => Err(Error::from(status)), 31 | }, 32 | None => Err(Error::Unexpected(format!( 33 | "unknown error code {}", 34 | status.error_code 35 | ))), 36 | } 37 | } 38 | -------------------------------------------------------------------------------- /src/value.rs: -------------------------------------------------------------------------------- 1 | use std::borrow::Cow; 2 | 3 | use crate::proto::{ 4 | self, 5 | schema::{ 6 | field_data::Field, scalar_field::Data as ScalarData, vector_field::Data as VectorData, 7 | DataType, 8 | }, 9 | }; 10 | 11 | pub enum Value<'a> { 12 | None, 13 | Bool(bool), 14 | Int8(i8), 15 | Int16(i16), 16 | Int32(i32), 17 | Long(i64), 18 | Float(f32), 19 | Double(f64), 20 | FloatArray(Cow<'a, [f32]>), 21 | Binary(Cow<'a, [u8]>), 22 | String(Cow<'a, str>), 23 | Json(Cow<'a, [u8]>), 24 | Array(Cow<'a, proto::schema::ScalarField>), 25 | } 26 | 27 | macro_rules! impl_from_for_field_data_column { 28 | ( $($t: ty, $o: ident ),+ ) => {$( 29 | impl From<$t> for Value<'static> { 30 | fn from(v: $t) -> Self { 31 | Self::$o(v as _) 32 | } 33 | } 34 | )*}; 35 | } 36 | 37 | impl_from_for_field_data_column! { 38 | bool, Bool, 39 | i8, Int8, 40 | i16, Int16, 41 | i32, Int32, 42 | i64, Long, 43 | f32, Float, 44 | f64, Double 45 | } 46 | 47 | impl Value<'_> { 48 | pub fn data_type(&self) -> DataType { 49 | match self { 50 | Value::None => DataType::None, 51 | Value::Bool(_) => DataType::Bool, 52 | Value::Int8(_) => DataType::Int8, 53 | Value::Int16(_) => DataType::Int16, 54 | Value::Int32(_) => DataType::Int32, 55 | Value::Long(_) => DataType::Int64, 56 | Value::Float(_) => DataType::Float, 57 | Value::Double(_) => DataType::Double, 58 | Value::String(_) => DataType::String, 59 | Value::Json(_) => DataType::Json, 60 | Value::FloatArray(_) => DataType::FloatVector, 61 | Value::Binary(_) => DataType::BinaryVector, 62 | Value::Array(_) => DataType::Array, 63 | } 64 | } 65 | } 66 | 67 | impl From for Value<'static> { 68 | fn from(v: String) -> Self { 69 | Self::String(Cow::Owned(v)) 70 | } 71 | } 72 | 73 | impl<'a> From<&'a str> for Value<'a> { 74 | fn from(v: &'a str) -> Self { 75 | Self::String(Cow::Borrowed(v)) 76 | } 77 | } 78 | 79 | impl<'a> From<&'a [u8]> for Value<'a> { 80 | fn from(v: &'a [u8]) -> Self { 81 | Self::Binary(Cow::Borrowed(v)) 82 | } 83 | } 84 | 85 | impl From> for Value<'static> { 86 | fn from(v: Vec) -> Self { 87 | Self::Binary(Cow::Owned(v)) 88 | } 89 | } 90 | 91 | impl<'a> From<&'a [f32]> for Value<'a> { 92 | fn from(v: &'a [f32]) -> Self { 93 | Self::FloatArray(Cow::Borrowed(v)) 94 | } 95 | } 96 | 97 | impl From> for Value<'static> { 98 | fn from(v: Vec) -> Self { 99 | Self::FloatArray(Cow::Owned(v)) 100 | } 101 | } 102 | 103 | #[derive(Debug, Clone)] 104 | pub enum ValueVec { 105 | None, 106 | Bool(Vec), 107 | Int(Vec), 108 | Long(Vec), 109 | Float(Vec), 110 | Double(Vec), 111 | Binary(Vec), 112 | String(Vec), 113 | Json(Vec>), 114 | Array(Vec), 115 | } 116 | 117 | macro_rules! impl_from_for_value_vec { 118 | ( $($t: ty, $o: ident ),+ ) => {$( 119 | impl From<$t> for ValueVec { 120 | fn from(v: $t) -> Self { 121 | Self::$o(v) 122 | } 123 | } 124 | )*}; 125 | } 126 | 127 | impl_from_for_value_vec! { 128 | Vec, Bool, 129 | Vec, Int, 130 | Vec, Long, 131 | Vec, String, 132 | Vec, Binary, 133 | Vec, Float, 134 | Vec, Double 135 | } 136 | 137 | macro_rules! impl_try_from_for_value_vec { 138 | ( $($o: ident, $t: ty ),+ ) => {$( 139 | impl TryFrom for $t { 140 | type Error = crate::error::Error; 141 | fn try_from(value: ValueVec) -> Result { 142 | match value { 143 | ValueVec::$o(v) => Ok(v), 144 | _ => Err(crate::error::Error::Conversion), 145 | } 146 | } 147 | } 148 | )*}; 149 | } 150 | 151 | impl_try_from_for_value_vec! { 152 | Bool, Vec, 153 | Int, Vec, 154 | Long, Vec, 155 | String, Vec, 156 | Binary, Vec, 157 | Float, Vec, 158 | Double, Vec 159 | } 160 | 161 | impl From> for ValueVec { 162 | fn from(v: Vec) -> Self { 163 | Self::Int(v.into_iter().map(Into::into).collect()) 164 | } 165 | } 166 | 167 | impl From> for ValueVec { 168 | fn from(v: Vec) -> Self { 169 | Self::Int(v.into_iter().map(Into::into).collect()) 170 | } 171 | } 172 | 173 | impl ValueVec { 174 | pub fn new(dtype: DataType) -> Self { 175 | match dtype { 176 | DataType::None => Self::None, 177 | DataType::Bool => Self::Bool(Vec::new()), 178 | DataType::Int8 => Self::Int(Vec::new()), 179 | DataType::Int16 => Self::Int(Vec::new()), 180 | DataType::Int32 => Self::Int(Vec::new()), 181 | DataType::Int64 => Self::Long(Vec::new()), 182 | DataType::Float => Self::Float(Vec::new()), 183 | DataType::Double => Self::Double(Vec::new()), 184 | DataType::String => Self::String(Vec::new()), 185 | DataType::VarChar => Self::String(Vec::new()), 186 | DataType::Json => Self::String(Vec::new()), 187 | DataType::Array => Self::Array(Vec::new()), 188 | DataType::BinaryVector => Self::Binary(Vec::new()), 189 | DataType::FloatVector => Self::Float(Vec::new()), 190 | DataType::Float16Vector => Self::Binary(Vec::new()), 191 | DataType::BFloat16Vector => Self::Binary(Vec::new()), 192 | } 193 | } 194 | 195 | pub fn check_dtype(&self, dtype: DataType) -> bool { 196 | match (self, dtype) { 197 | (ValueVec::Binary(..), DataType::BinaryVector) 198 | | (ValueVec::Float(..), DataType::FloatVector) 199 | | (ValueVec::Float(..), DataType::Float) 200 | | (ValueVec::Int(..), DataType::Int8) 201 | | (ValueVec::Int(..), DataType::Int16) 202 | | (ValueVec::Int(..), DataType::Int32) 203 | | (ValueVec::Long(..), DataType::Int64) 204 | | (ValueVec::Bool(..), DataType::Bool) 205 | | (ValueVec::String(..), DataType::String) 206 | | (ValueVec::String(..), DataType::VarChar) 207 | | (ValueVec::None, _) 208 | | (ValueVec::Double(..), DataType::Double) => true, 209 | _ => false, 210 | } 211 | } 212 | 213 | #[inline] 214 | pub fn is_empty(&self) -> bool { 215 | self.len() == 0 216 | } 217 | 218 | pub fn len(&self) -> usize { 219 | match self { 220 | ValueVec::None => 0, 221 | ValueVec::Bool(v) => v.len(), 222 | ValueVec::Int(v) => v.len(), 223 | ValueVec::Long(v) => v.len(), 224 | ValueVec::Float(v) => v.len(), 225 | ValueVec::Double(v) => v.len(), 226 | ValueVec::Binary(v) => v.len(), 227 | ValueVec::String(v) => v.len(), 228 | ValueVec::Json(v) => v.len(), 229 | ValueVec::Array(v) => v.len(), 230 | } 231 | } 232 | 233 | pub fn clear(&mut self) { 234 | match self { 235 | ValueVec::None => (), 236 | ValueVec::Bool(v) => v.clear(), 237 | ValueVec::Int(v) => v.clear(), 238 | ValueVec::Long(v) => v.clear(), 239 | ValueVec::Float(v) => v.clear(), 240 | ValueVec::Double(v) => v.clear(), 241 | ValueVec::Binary(v) => v.clear(), 242 | ValueVec::String(v) => v.clear(), 243 | ValueVec::Json(v) => v.clear(), 244 | ValueVec::Array(v) => v.clear(), 245 | } 246 | } 247 | } 248 | 249 | impl From for ValueVec { 250 | fn from(f: Field) -> Self { 251 | match f { 252 | Field::Scalars(s) => match s.data { 253 | Some(x) => match x { 254 | ScalarData::BoolData(v) => Self::Bool(v.data), 255 | ScalarData::IntData(v) => Self::Int(v.data), 256 | ScalarData::LongData(v) => Self::Long(v.data), 257 | ScalarData::FloatData(v) => Self::Float(v.data), 258 | ScalarData::DoubleData(v) => Self::Double(v.data), 259 | ScalarData::StringData(v) => Self::String(v.data), 260 | ScalarData::JsonData(v) => Self::Json(v.data), 261 | ScalarData::ArrayData(v) => Self::Array(v.data), 262 | ScalarData::BytesData(_) => unimplemented!(), // Self::Bytes(v.data), 263 | }, 264 | None => Self::None, 265 | }, 266 | 267 | Field::Vectors(arr) => match arr.data { 268 | Some(x) => match x { 269 | VectorData::FloatVector(v) => Self::Float(v.data), 270 | VectorData::BinaryVector(v) => Self::Binary(v), 271 | VectorData::Bfloat16Vector(v) => Self::Binary(v), 272 | VectorData::Float16Vector(v) => Self::Binary(v), 273 | }, 274 | None => Self::None, 275 | }, 276 | } 277 | } 278 | } 279 | 280 | macro_rules! impl_try_from_for_value_column { 281 | ( $($o: ident,$t: ty ),+ ) => {$( 282 | impl TryFrom> for $t { 283 | type Error = crate::error::Error; 284 | fn try_from(value: Value<'_>) -> Result { 285 | match value { 286 | Value::$o(v) => Ok(v), 287 | _ => Err(crate::error::Error::Conversion), 288 | } 289 | } 290 | } 291 | )*}; 292 | } 293 | 294 | impl_try_from_for_value_column! { 295 | Bool,bool, 296 | Int8,i8, 297 | Int16,i16, 298 | Int32,i32, 299 | Long,i64, 300 | Float,f32, 301 | Double,f64 302 | } 303 | 304 | #[cfg(test)] 305 | mod test { 306 | use crate::{ 307 | error::Error, 308 | value::{Value, ValueVec}, 309 | }; 310 | 311 | #[test] 312 | fn test_try_from_for_value_column() { 313 | // bool 314 | let b = Value::Bool(false); 315 | let b: Result = b.try_into(); 316 | assert!(b.is_ok()); 317 | assert!(!b.unwrap()); 318 | //i8 319 | let int8 = Value::Int8(12); 320 | let r: Result = int8.try_into(); 321 | assert!(r.is_ok()); 322 | assert_eq!(12, r.unwrap()); 323 | //i16 324 | let int16 = Value::Int16(1225); 325 | let r: Result = int16.try_into(); 326 | assert!(r.is_ok()); 327 | assert_eq!(1225, r.unwrap()); 328 | //i32 329 | let int32 = Value::Int32(37360798); 330 | let r: Result = int32.try_into(); 331 | assert!(r.is_ok()); 332 | assert_eq!(37360798, r.unwrap()); 333 | //i64 334 | let long = Value::Long(123); 335 | let r: Result = long.try_into(); 336 | assert!(r.is_ok()); 337 | assert_eq!(123, r.unwrap()); 338 | 339 | let float = Value::Float(22104f32); 340 | let r: Result = float.try_into(); 341 | assert!(r.is_ok()); 342 | assert_eq!(22104f32, r.unwrap()); 343 | 344 | let double = Value::Double(22104f64); 345 | let r: Result = double.try_into(); 346 | assert!(r.is_ok()); 347 | assert_eq!(22104f64, r.unwrap()); 348 | } 349 | 350 | #[test] 351 | fn test_try_from_for_value_vec() { 352 | let b = ValueVec::Bool(vec![false, false]); 353 | let b: Result, Error> = b.try_into(); 354 | assert!(b.is_ok()); 355 | assert_eq!(vec![false, false], b.unwrap()); 356 | 357 | let b = ValueVec::Int(vec![1, 2]); 358 | let b: Result, Error> = b.try_into(); 359 | assert!(b.is_ok()); 360 | assert_eq!(vec![1, 2], b.unwrap()); 361 | 362 | let v: Vec = vec![4095291003, 2581116377, 3892395808]; 363 | let b = ValueVec::Long(v.clone()); 364 | let b: Result, Error> = b.try_into(); 365 | assert!(b.is_ok()); 366 | assert_eq!(v, b.unwrap()); 367 | 368 | let v: Vec = vec![11., 7., 754., 68., 34.]; 369 | let b = ValueVec::Float(v.clone()); 370 | let b: Result, Error> = b.try_into(); 371 | assert!(b.is_ok()); 372 | assert_eq!(v, b.unwrap()); 373 | 374 | let v: Vec = vec![28., 9., 92., 6099786761., 64.]; 375 | let b = ValueVec::Double(v.clone()); 376 | let b_err: Result, Error> = b.clone().try_into(); 377 | assert!(b_err.is_err()); 378 | let b: Result, Error> = b.try_into(); 379 | assert_eq!(v, b.unwrap()); 380 | 381 | let v = vec![28, 5, 70, 62, 57, 103, 68]; 382 | let b = ValueVec::Binary(v.clone()); 383 | let b: Result, Error> = b.try_into(); 384 | assert!(b.is_ok()); 385 | assert_eq!(v, b.unwrap()); 386 | 387 | let v: Vec = vec!["Janoato", "Samoa", "opde@tuwuv.yt"] 388 | .into_iter() 389 | .map(|a| a.into()) 390 | .collect(); 391 | let b = ValueVec::String(v.clone()); 392 | let b: Result, Error> = b.try_into(); 393 | assert!(b.is_ok()); 394 | assert_eq!(v, b.unwrap()); 395 | } 396 | } 397 | -------------------------------------------------------------------------------- /tests/alias.rs: -------------------------------------------------------------------------------- 1 | mod common; 2 | use common::*; 3 | use milvus::{client::Client, error::Result}; 4 | 5 | async fn clean_test_collection(client: Client, collection_name: &str) -> Result<()> { 6 | client.drop_collection(collection_name).await?; 7 | Ok(()) 8 | } 9 | 10 | #[tokio::test] 11 | async fn test_create_alias() -> Result<()> { 12 | let alias = "test_create_alias"; 13 | let (client, schema) = create_test_collection(true).await?; 14 | client.create_alias(schema.name(), alias).await?; 15 | client.drop_alias(alias).await?; 16 | clean_test_collection(client, schema.name()).await?; 17 | Ok(()) 18 | } 19 | 20 | #[tokio::test] 21 | async fn test_alter_alias() -> Result<()> { 22 | let alias = "test_alter_alias"; 23 | let (client1, schema1) = create_test_collection(true).await?; 24 | client1.create_alias(schema1.name(), alias).await?; 25 | let (client2, schema2) = create_test_collection(true).await?; 26 | client2.alter_alias(schema2.name(), alias).await?; 27 | client2.drop_alias(alias).await?; 28 | clean_test_collection(client1, schema1.name()).await?; 29 | clean_test_collection(client2, schema2.name()).await?; 30 | Ok(()) 31 | } 32 | -------------------------------------------------------------------------------- /tests/client.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use milvus::client::*; 18 | use milvus::error::Result; 19 | use milvus::options::CreateCollectionOptions; 20 | use milvus::schema::*; 21 | 22 | mod common; 23 | use common::*; 24 | 25 | #[tokio::test] 26 | async fn create_client() -> Result<()> { 27 | match Client::new(URL).await { 28 | Ok(_) => Result::<()>::Ok(()), 29 | Err(e) => panic!("Error is {}.", e), 30 | } 31 | } 32 | 33 | #[tokio::test] 34 | async fn create_client_wrong_url() -> Result<()> { 35 | const URL: &str = "http://localhost:9999"; 36 | match Client::new(URL).await { 37 | Ok(_) => panic!("Should fail due to wrong url."), 38 | Err(_) => Result::<()>::Ok(()), 39 | } 40 | } 41 | 42 | #[tokio::test] 43 | async fn create_client_wrong_fmt() -> Result<()> { 44 | const URL: &str = "9999"; 45 | match Client::new(URL).await { 46 | Ok(_) => panic!("Should fail due to wrong format url."), 47 | Err(_) => Result::<()>::Ok(()), 48 | } 49 | } 50 | 51 | #[tokio::test] 52 | async fn has_collection() -> Result<()> { 53 | const NAME: &str = "qwerty"; 54 | let client = Client::new(URL).await?; 55 | match client.has_collection(NAME).await { 56 | Ok(has) => { 57 | if has { 58 | panic!("Expect no such collection."); 59 | } else { 60 | Ok(()) 61 | } 62 | } 63 | Err(e) => Err(e), 64 | } 65 | } 66 | 67 | #[tokio::test] 68 | async fn create_has_drop_collection() -> Result<()> { 69 | const NAME: &str = "create_has_drop_collection"; 70 | 71 | let client = Client::new(URL).await?; 72 | // let client = ClientBuilder::new(URL).username("username").password("password").build().await?; 73 | 74 | let mut schema = CollectionSchemaBuilder::new(NAME, "hello world"); 75 | let schema = schema 76 | .add_field(FieldSchema::new_int64("i64_field", "")) 77 | .add_field(FieldSchema::new_bool("bool_field", "")) 78 | .set_primary_key("i64_field")? 79 | .enable_auto_id()? 80 | .build()?; 81 | 82 | if client.has_collection(NAME).await? { 83 | client.drop_collection(NAME).await?; 84 | } 85 | 86 | let collection = client 87 | .create_collection( 88 | schema, 89 | Some(CreateCollectionOptions::with_consistency_level( 90 | ConsistencyLevel::Session, 91 | )), 92 | ) 93 | .await?; 94 | 95 | assert!(client.has_collection(NAME).await?); 96 | 97 | client.drop_collection(NAME).await?; 98 | assert!(!client.has_collection(NAME).await?); 99 | 100 | Ok(()) 101 | } 102 | 103 | #[tokio::test] 104 | async fn create_alter_drop_alias() -> Result<()> { 105 | let alias0 = gen_random_name(); 106 | let alias1 = gen_random_name(); 107 | 108 | let client = Client::new(URL).await?; 109 | 110 | let (_, schema1) = create_test_collection(true).await?; 111 | let (_, schema2) = create_test_collection(true).await?; 112 | 113 | client.create_alias(schema1.name(), &alias0).await?; 114 | assert!(client.has_collection(alias0).await?); 115 | 116 | client.create_alias(schema2.name(), &alias1).await?; 117 | 118 | client.alter_alias(schema1.name(), &alias1).await?; 119 | 120 | client.drop_collection(schema2.name()).await?; 121 | assert!(client.has_collection(alias1).await?); 122 | 123 | Ok(()) 124 | } 125 | -------------------------------------------------------------------------------- /tests/collection.rs: -------------------------------------------------------------------------------- 1 | // Licensed to the LF AI & Data foundation under one 2 | // or more contributor license agreements. See the NOTICE file 3 | // distributed with this work for additional information 4 | // regarding copyright ownership. The ASF licenses this file 5 | // to you under the Apache License, Version 2.0 (the 6 | // "License"); you may not use this file except in compliance 7 | // with the License. You may obtain a copy of the License at 8 | // 9 | // http://www.apache.org/licenses/LICENSE-2.0 10 | // 11 | // Unless required by applicable law or agreed to in writing, software 12 | // distributed under the License is distributed on an "AS IS" BASIS, 13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | // See the License for the specific language governing permissions and 15 | // limitations under the License. 16 | 17 | use milvus::client::ConsistencyLevel; 18 | use milvus::collection::{Collection, ParamValue}; 19 | use milvus::data::FieldColumn; 20 | use milvus::error::Result; 21 | use milvus::index::{IndexParams, IndexType, MetricType}; 22 | use milvus::mutate::InsertOptions; 23 | use milvus::options::LoadOptions; 24 | use milvus::query::{QueryOptions, SearchOptions}; 25 | use std::collections::HashMap; 26 | 27 | mod common; 28 | use common::*; 29 | 30 | use milvus::value::ValueVec; 31 | 32 | #[tokio::test] 33 | async fn manual_compaction_empty_collection() -> Result<()> { 34 | let (client, schema) = create_test_collection(true).await?; 35 | let resp = client.manual_compaction(schema.name()).await?; 36 | assert_eq!(0, resp.plan_count); 37 | Ok(()) 38 | } 39 | 40 | #[tokio::test] 41 | async fn collection_upsert() -> Result<()> { 42 | let (client, schema) = create_test_collection(false).await?; 43 | let pk_data = gen_random_int64_vector(2000); 44 | let vec_data = gen_random_f32_vector(DEFAULT_DIM * 2000); 45 | let pk_col = FieldColumn::new(schema.get_field("id").unwrap(), pk_data); 46 | let vec_col = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), vec_data); 47 | client 48 | .upsert(schema.name(), vec![pk_col, vec_col], None) 49 | .await?; 50 | let index_params = IndexParams::new( 51 | DEFAULT_INDEX_NAME.to_owned(), 52 | IndexType::IvfFlat, 53 | milvus::index::MetricType::L2, 54 | HashMap::from([("nlist".to_owned(), "32".to_owned())]), 55 | ); 56 | client 57 | .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) 58 | .await?; 59 | client 60 | .load_collection(schema.name(), Some(LoadOptions::default())) 61 | .await?; 62 | 63 | let options = QueryOptions::default(); 64 | let options = options.output_fields(vec![String::from("count(*)")]); 65 | let result = client.query(schema.name(), "", &options).await?; 66 | if let ValueVec::Long(vec) = &result[0].value { 67 | assert_eq!(2000, vec[0]); 68 | } else { 69 | panic!("invalid result"); 70 | } 71 | Ok(()) 72 | } 73 | 74 | #[tokio::test] 75 | async fn collection_basic() -> Result<()> { 76 | let (client, schema) = create_test_collection(true).await?; 77 | 78 | let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); 79 | 80 | let embed_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); 81 | 82 | client 83 | .insert(schema.name(), vec![embed_column], None) 84 | .await?; 85 | client.flush(schema.name()).await?; 86 | let index_params = IndexParams::new( 87 | DEFAULT_INDEX_NAME.to_owned(), 88 | IndexType::IvfFlat, 89 | milvus::index::MetricType::L2, 90 | HashMap::from([("nlist".to_owned(), "32".to_owned())]), 91 | ); 92 | client 93 | .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) 94 | .await?; 95 | client 96 | .load_collection(schema.name(), Some(LoadOptions::default())) 97 | .await?; 98 | 99 | let options = QueryOptions::default(); 100 | let result = client.query(schema.name(), "id > 0", &options).await?; 101 | 102 | println!( 103 | "result num: {}", 104 | result.first().map(|c| c.len()).unwrap_or(0), 105 | ); 106 | 107 | client.drop_collection(schema.name()).await?; 108 | Ok(()) 109 | } 110 | 111 | #[tokio::test] 112 | async fn collection_index() -> Result<()> { 113 | let (client, schema) = create_test_collection(true).await?; 114 | 115 | let feature = gen_random_f32_vector(DEFAULT_DIM * 2000); 116 | 117 | let feature_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), feature); 118 | 119 | client 120 | .insert(schema.name(), vec![feature_column], None) 121 | .await?; 122 | client.flush(schema.name()).await?; 123 | 124 | let index_params = IndexParams::new( 125 | DEFAULT_INDEX_NAME.to_owned(), 126 | IndexType::IvfFlat, 127 | milvus::index::MetricType::L2, 128 | HashMap::from([("nlist".to_owned(), "32".to_owned())]), 129 | ); 130 | client 131 | .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params.clone()) 132 | .await?; 133 | let index_list = client 134 | .describe_index(schema.name(), DEFAULT_VEC_FIELD) 135 | .await?; 136 | assert!(index_list.len() == 1, "{}", index_list.len()); 137 | let index = &index_list[0]; 138 | 139 | assert_eq!(index.params().name(), index_params.name()); 140 | assert_eq!(index.params().extra_params(), index_params.extra_params()); 141 | 142 | client.drop_index(schema.name(), DEFAULT_VEC_FIELD).await?; 143 | client.drop_collection(schema.name()).await?; 144 | Ok(()) 145 | } 146 | 147 | #[tokio::test] 148 | async fn collection_search() -> Result<()> { 149 | let (client, schema) = create_test_collection(true).await?; 150 | 151 | let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); 152 | let embed_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); 153 | 154 | client 155 | .insert(schema.name(), vec![embed_column], None) 156 | .await?; 157 | client.flush(schema.name()).await?; 158 | let index_params = IndexParams::new( 159 | "ivf_flat".to_owned(), 160 | IndexType::IvfFlat, 161 | MetricType::L2, 162 | HashMap::from_iter([("nlist".to_owned(), 32.to_string())]), 163 | ); 164 | client 165 | .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) 166 | .await?; 167 | client.flush(schema.name()).await?; 168 | client 169 | .load_collection(schema.name(), Some(LoadOptions::default())) 170 | .await?; 171 | 172 | let mut option = SearchOptions::with_limit(10) 173 | .metric_type(MetricType::L2) 174 | .output_fields(vec!["id".to_owned()]); 175 | option = option.add_param("nprobe", ParamValue!(16)); 176 | let query_vec = gen_random_f32_vector(DEFAULT_DIM); 177 | 178 | let result = client 179 | .search( 180 | schema.name(), 181 | vec![query_vec.into()], 182 | DEFAULT_VEC_FIELD, 183 | &option, 184 | ) 185 | .await?; 186 | 187 | assert_eq!(result[0].size, 10); 188 | 189 | client.drop_collection(schema.name()).await?; 190 | Ok(()) 191 | } 192 | 193 | #[tokio::test] 194 | async fn collection_range_search() -> Result<()> { 195 | let (client, schema) = create_test_collection(true).await?; 196 | 197 | let embed_data = gen_random_f32_vector(DEFAULT_DIM * 2000); 198 | let embed_column = FieldColumn::new(schema.get_field(DEFAULT_VEC_FIELD).unwrap(), embed_data); 199 | 200 | client 201 | .insert(schema.name(), vec![embed_column], None) 202 | .await?; 203 | client.flush(schema.name()).await?; 204 | let index_params = IndexParams::new( 205 | "ivf_flat".to_owned(), 206 | IndexType::IvfFlat, 207 | MetricType::L2, 208 | HashMap::from_iter([("nlist".to_owned(), 32.to_string())]), 209 | ); 210 | client 211 | .create_index(schema.name(), DEFAULT_VEC_FIELD, index_params) 212 | .await?; 213 | client.flush(schema.name()).await?; 214 | client 215 | .load_collection(schema.name(), Some(LoadOptions::default())) 216 | .await?; 217 | 218 | let radius_limit: f32 = 20.0; 219 | let range_filter_limit: f32 = 10.0; 220 | 221 | let mut option = SearchOptions::with_limit(5) 222 | .metric_type(MetricType::L2) 223 | .output_fields(vec!["id".to_owned()]); 224 | option = option.add_param("nprobe", ParamValue!(16)); 225 | option = option.radius(radius_limit).range_filter(range_filter_limit); 226 | let query_vec = gen_random_f32_vector(DEFAULT_DIM); 227 | 228 | let result = client 229 | .search( 230 | schema.name(), 231 | vec![query_vec.into()], 232 | DEFAULT_VEC_FIELD, 233 | &option, 234 | ) 235 | .await?; 236 | 237 | for record in &result { 238 | for value in &record.score { 239 | assert!(*value >= range_filter_limit && *value <= radius_limit); 240 | } 241 | } 242 | 243 | client.drop_collection(schema.name()).await?; 244 | Ok(()) 245 | } 246 | -------------------------------------------------------------------------------- /tests/common/mod.rs: -------------------------------------------------------------------------------- 1 | use milvus::client::*; 2 | use milvus::error::Result; 3 | use milvus::options::CreateCollectionOptions; 4 | use milvus::schema::{CollectionSchema, CollectionSchemaBuilder, FieldSchema}; 5 | use rand::Rng; 6 | 7 | pub const DEFAULT_DIM: i64 = 128; 8 | pub const DEFAULT_VEC_FIELD: &str = "feature"; 9 | pub const DEFAULT_INDEX_NAME: &str = "feature_index"; 10 | pub const URL: &str = "http://localhost:19530"; 11 | 12 | pub async fn create_test_collection(autoid: bool) -> Result<(Client, CollectionSchema)> { 13 | let collection_name = gen_random_name(); 14 | let collection_name = format!("{}_{}", "test_collection", collection_name); 15 | let client = Client::new(URL).await?; 16 | let schema = CollectionSchemaBuilder::new(&collection_name, "") 17 | .add_field(FieldSchema::new_primary_int64("id", "", autoid)) 18 | .add_field(FieldSchema::new_float_vector( 19 | DEFAULT_VEC_FIELD, 20 | "", 21 | DEFAULT_DIM, 22 | )) 23 | .build()?; 24 | if client.has_collection(&collection_name).await? { 25 | client.drop_collection(&collection_name).await?; 26 | } 27 | client 28 | .create_collection( 29 | schema.clone(), 30 | Some(CreateCollectionOptions::with_consistency_level( 31 | ConsistencyLevel::Eventually, 32 | )), 33 | ) 34 | .await?; 35 | Ok((client, schema)) 36 | } 37 | 38 | pub fn gen_random_name() -> String { 39 | format!( 40 | "r{}", 41 | rand::thread_rng() 42 | .sample_iter(&rand::distributions::Alphanumeric) 43 | .take(7) 44 | .map(char::from) 45 | .collect::(), 46 | ) 47 | } 48 | 49 | pub fn gen_random_int64_vector(n: i64) -> Vec { 50 | let mut data: Vec = Vec::with_capacity(n as usize); 51 | let mut rng = rand::thread_rng(); 52 | for _ in 0..n { 53 | data.push(rng.gen()); 54 | } 55 | data 56 | } 57 | 58 | pub fn gen_random_f32_vector(n: i64) -> Vec { 59 | let mut data = Vec::::with_capacity(n as usize); 60 | let mut rng = rand::thread_rng(); 61 | for _ in 0..n { 62 | data.push(rng.gen()); 63 | } 64 | data 65 | } 66 | --------------------------------------------------------------------------------