llama-index-vector-stores-bigquery 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,153 @@
1
+ llama_index/_static
2
+ .DS_Store
3
+ # Byte-compiled / optimized / DLL files
4
+ __pycache__/
5
+ *.py[cod]
6
+ *$py.class
7
+
8
+ # C extensions
9
+ *.so
10
+
11
+ # Distribution / packaging
12
+ .Python
13
+ bin/
14
+ build/
15
+ develop-eggs/
16
+ dist/
17
+ downloads/
18
+ eggs/
19
+ .eggs/
20
+ etc/
21
+ include/
22
+ lib/
23
+ lib64/
24
+ parts/
25
+ sdist/
26
+ share/
27
+ var/
28
+ wheels/
29
+ pip-wheel-metadata/
30
+ share/python-wheels/
31
+ *.egg-info/
32
+ .installed.cfg
33
+ *.egg
34
+ MANIFEST
35
+
36
+ # PyInstaller
37
+ # Usually these files are written by a python script from a template
38
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
39
+ *.manifest
40
+ *.spec
41
+
42
+ # Installer logs
43
+ pip-log.txt
44
+ pip-delete-this-directory.txt
45
+
46
+ # Unit test / coverage reports
47
+ htmlcov/
48
+ .tox/
49
+ .nox/
50
+ .coverage
51
+ .coverage.*
52
+ .cache
53
+ nosetests.xml
54
+ coverage.xml
55
+ *.cover
56
+ *.py,cover
57
+ .hypothesis/
58
+ .pytest_cache/
59
+ .ruff_cache
60
+
61
+ # Translations
62
+ *.mo
63
+ *.pot
64
+
65
+ # Django stuff:
66
+ *.log
67
+ local_settings.py
68
+ db.sqlite3
69
+ db.sqlite3-journal
70
+
71
+ # Flask stuff:
72
+ instance/
73
+ .webassets-cache
74
+
75
+ # Scrapy stuff:
76
+ .scrapy
77
+
78
+ # Sphinx documentation
79
+ docs/_build/
80
+
81
+ # PyBuilder
82
+ target/
83
+
84
+ # Jupyter Notebook
85
+ .ipynb_checkpoints
86
+ notebooks/
87
+
88
+ # IPython
89
+ profile_default/
90
+ ipython_config.py
91
+
92
+ # pyenv
93
+ .python-version
94
+
95
+ # pipenv
96
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
97
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
98
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
99
+ # install all needed dependencies.
100
+ #Pipfile.lock
101
+
102
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow
103
+ __pypackages__/
104
+
105
+ # Celery stuff
106
+ celerybeat-schedule
107
+ celerybeat.pid
108
+
109
+ # SageMath parsed files
110
+ *.sage.py
111
+
112
+ # Environments
113
+ .env
114
+ .venv
115
+ env/
116
+ venv/
117
+ ENV/
118
+ env.bak/
119
+ venv.bak/
120
+ pyvenv.cfg
121
+
122
+ # Spyder project settings
123
+ .spyderproject
124
+ .spyproject
125
+
126
+ # Rope project settings
127
+ .ropeproject
128
+
129
+ # mkdocs documentation
130
+ /site
131
+
132
+ # mypy
133
+ .mypy_cache/
134
+ .dmypy.json
135
+ dmypy.json
136
+
137
+ # Pyre type checker
138
+ .pyre/
139
+
140
+ # Jetbrains
141
+ .idea
142
+ modules/
143
+ *.swp
144
+
145
+ # VsCode
146
+ .vscode
147
+
148
+ # pipenv
149
+ Pipfile
150
+ Pipfile.lock
151
+
152
+ # pyright
153
+ pyrightconfig.json
@@ -0,0 +1,21 @@
1
+ The MIT License
2
+
3
+ Copyright (c) Jerry Liu
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in
13
+ all copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
21
+ THE SOFTWARE.
@@ -0,0 +1,44 @@
1
+ Metadata-Version: 2.4
2
+ Name: llama-index-vector-stores-bigquery
3
+ Version: 0.1.0
4
+ Summary: llama-index vector_stores bigquery vector search integration
5
+ License-Expression: MIT
6
+ License-File: LICENSE
7
+ Requires-Python: <4.0,>=3.9
8
+ Requires-Dist: google-cloud-bigquery<4,>=3.34.0
9
+ Requires-Dist: llama-index-core<0.14,>=0.13.0
10
+ Description-Content-Type: text/markdown
11
+
12
+ # LlamaIndex Vector_Stores Integration: BigQuery
13
+
14
+ Vector store index using Google BigQuery.
15
+ Supports efficient storage and querying of vector embeddings using BigQuery's native vector search capabilities.
16
+ For more information, see the official [BigQuery Vector Search Documentation](https://cloud.google.com/bigquery/docs/vector-search-intro)
17
+
18
+ ## 🔐 Required IAM Permissions
19
+
20
+ To use this integration, ensure your account has the following permissions:
21
+
22
+ - `roles/bigquery.dataOwner` (BigQuery Data Owner)
23
+ - `roles/bigquery.dataEditor` (BigQuery Data Editor)
24
+
25
+ ## 🔧 Installation
26
+
27
+ ```bash
28
+ pip install llama-index-vector-stores-bigquery
29
+ ```
30
+
31
+ ## 💻 Example Usage
32
+
33
+ ```python
34
+ from google.cloud.bigquery import Client
35
+ from llama_index.vector_stores.bigquery import BigQueryVectorStore
36
+
37
+ client = Client()
38
+
39
+ vector_store = BigQueryVectorStore(
40
+ table_id="my_bigquery_table",
41
+ dataset_id="my_bigquery_dataset",
42
+ bigquery_client=client,
43
+ )
44
+ ```
@@ -0,0 +1,33 @@
1
+ # LlamaIndex Vector_Stores Integration: BigQuery
2
+
3
+ Vector store index using Google BigQuery.
4
+ Supports efficient storage and querying of vector embeddings using BigQuery's native vector search capabilities.
5
+ For more information, see the official [BigQuery Vector Search Documentation](https://cloud.google.com/bigquery/docs/vector-search-intro)
6
+
7
+ ## 🔐 Required IAM Permissions
8
+
9
+ To use this integration, ensure your account has the following permissions:
10
+
11
+ - `roles/bigquery.dataOwner` (BigQuery Data Owner)
12
+ - `roles/bigquery.dataEditor` (BigQuery Data Editor)
13
+
14
+ ## 🔧 Installation
15
+
16
+ ```bash
17
+ pip install llama-index-vector-stores-bigquery
18
+ ```
19
+
20
+ ## 💻 Example Usage
21
+
22
+ ```python
23
+ from google.cloud.bigquery import Client
24
+ from llama_index.vector_stores.bigquery import BigQueryVectorStore
25
+
26
+ client = Client()
27
+
28
+ vector_store = BigQueryVectorStore(
29
+ table_id="my_bigquery_table",
30
+ dataset_id="my_bigquery_dataset",
31
+ bigquery_client=client,
32
+ )
33
+ ```
@@ -0,0 +1,3 @@
1
+ from llama_index.vector_stores.bigquery.base import BigQueryVectorStore
2
+
3
+ __all__ = ["BigQueryVectorStore"]
@@ -0,0 +1,547 @@
1
+ """
2
+ Google BigQuery Vector Search.
3
+
4
+ BigQuery Vector Search is a fully managed feature of BigQuery that enables fast,
5
+ scalable similarity search over high-dimensional embeddings using approximate
6
+ nearest neighbor methods. For more information visit:
7
+ https://cloud.google.com/bigquery/docs/vector-search-intro
8
+ """
9
+
10
+ import logging
11
+ from enum import Enum
12
+ from typing import Any, Dict, List, Optional, Union
13
+
14
+ from google.auth import credentials
15
+ from google.cloud import bigquery
16
+ from llama_index.core.bridge.pydantic import PrivateAttr
17
+ from llama_index.core.schema import BaseNode, MetadataMode, TextNode
18
+ from llama_index.core.vector_stores.types import (
19
+ BasePydanticVectorStore,
20
+ MetadataFilters,
21
+ VectorStoreQuery,
22
+ VectorStoreQueryResult,
23
+ )
24
+ from llama_index.core.vector_stores.utils import (
25
+ metadata_dict_to_node,
26
+ node_to_metadata_dict,
27
+ )
28
+ from pydantic import BaseModel
29
+
30
+ from llama_index.vector_stores.bigquery.utils import build_where_clause_and_params
31
+
32
+
33
+ _logger = logging.getLogger(__name__)
34
+
35
+
36
+ class _BigQueryRow(BaseModel):
37
+ node_id: str
38
+ embedding: List[float]
39
+ text: str
40
+ metadata: Dict[str, Any]
41
+ distance: Optional[float] = None
42
+
43
+
44
+ class DistanceType(str, Enum):
45
+ EUCLIDEAN = "EUCLIDEAN"
46
+ COSINE = "COSINE"
47
+ DOT_PRODUCT = "DOT_PRODUCT"
48
+
49
+
50
+ class BigQueryVectorStore(BasePydanticVectorStore):
51
+ """
52
+ Vector store index using Google BigQuery.
53
+
54
+ Provides integration with BigQuery for storing and querying vector embeddings.
55
+ For more information, visit: https://cloud.google.com/bigquery/docs/vector-search-intro
56
+
57
+ Required IAM Permissions:
58
+ - `roles/bigquery.dataOwner` (BigQuery Data Owner)
59
+ - `roles/bigquery.dataEditor` (BigQuery Data Editor)
60
+
61
+ Examples:
62
+ `pip install llama-index-vector-stores-bigquery`
63
+
64
+ ```python
65
+ from google.cloud.bigquery import Client
66
+ from llama_index.vector_stores.bigquery import BigQueryVectorStore
67
+
68
+ client = Client()
69
+
70
+ vector_store = BigQueryVectorStore(
71
+ table_id="my_bigquery_table",
72
+ dataset_id="my_bigquery_dataset",
73
+ bigquery_client=client,
74
+ )
75
+ ```
76
+
77
+ """
78
+
79
+ stores_text: bool = True
80
+ distance_type: DistanceType = DistanceType.EUCLIDEAN
81
+
82
+ _table: bigquery.Table = PrivateAttr()
83
+ _dataset: bigquery.Dataset = PrivateAttr()
84
+ _client: bigquery.Client = PrivateAttr()
85
+ _full_table_id: str = PrivateAttr()
86
+
87
+ def __init__(
88
+ self,
89
+ table_id: str,
90
+ dataset_id: str,
91
+ project_id: Optional[str] = None,
92
+ region: Optional[str] = None,
93
+ distance_type: Optional[DistanceType] = DistanceType.EUCLIDEAN,
94
+ auth_credentials: Optional[credentials.Credentials] = None,
95
+ bigquery_client: Optional[bigquery.Client] = None,
96
+ **kwargs: Any,
97
+ ):
98
+ """
99
+ Initialize a BigQuery Vector store.
100
+
101
+ If a `bigquery_client` is provided, it will be used directly. Otherwise, a client will be initialized using
102
+ the optional `project_id`, `region`, and/or `auth_credentials`. If none are provided, default credentials
103
+ will be used. For details on authentication, visit:
104
+ https://googleapis.dev/python/google-api-core/latest/auth.html
105
+
106
+ Args:
107
+ table_id: The ID of the BigQuery table to use for vector storage.
108
+ dataset_id: The ID of the dataset containing the table.
109
+ project_id: The GCP project ID. If not provided, it will be inferred from the client or environment.
110
+ region: Optionally specify a default location for datasets / tables.
111
+ distance_type: Optionally specify a distance type to use `EUCLIDEAN`, `COSINE`, or `DOT_PRODUCT`.
112
+ auth_credentials: Optional credentials object used to authenticate with BigQuery.
113
+ bigquery_client: An existing BigQuery client instance. If not provided, one will be created.
114
+ **kwargs: Additional keyword arguments passed to the parent class.
115
+
116
+ """
117
+ super().__init__(
118
+ **kwargs,
119
+ )
120
+
121
+ self._client: bigquery.Client = bigquery_client or self._initialize_client(
122
+ project_id, region, auth_credentials
123
+ )
124
+ self._dataset: bigquery.Dataset = self._create_dataset_if_not_exists(dataset_id)
125
+ self._table: bigquery.Table = self._create_table_if_not_exists(table_id)
126
+ self._full_table_id: str = (
127
+ f"{self._client.project}.{self._dataset.dataset_id}.{self._table.table_id}"
128
+ )
129
+ self.distance_type: DistanceType = DistanceType(distance_type)
130
+
131
+ @classmethod
132
+ def from_params(
133
+ cls,
134
+ table_id: str,
135
+ dataset_id: str,
136
+ project_id: Optional[str] = None,
137
+ region: Optional[str] = None,
138
+ distance_type: Optional[DistanceType] = DistanceType.EUCLIDEAN,
139
+ auth_credentials: Optional[credentials.Credentials] = None,
140
+ bigquery_client: Optional[bigquery.Client] = None,
141
+ ) -> "BigQueryVectorStore":
142
+ """
143
+ Initialize a BigQuery Vector store.
144
+
145
+ Args:
146
+ table_id: The ID of the BigQuery table to use for vector storage.
147
+ dataset_id: The ID of the dataset containing the table.
148
+ project_id: The GCP project ID. If not provided, it will be inferred from the client or environment.
149
+ region: Optionally specify a default location for datasets / tables.
150
+ distance_type: Optionally specify a distance type to use `EUCLIDEAN`, `COSINE`, or `DOT_PRODUCT`.
151
+ auth_credentials: Optional credentials object used to authenticate with BigQuery.
152
+ bigquery_client: An existing BigQuery client instance. If not provided, one will be created.
153
+
154
+ Returns:
155
+ BigQueryVectorStore
156
+
157
+ """
158
+ return cls(
159
+ table_id=table_id,
160
+ dataset_id=dataset_id,
161
+ project_id=project_id,
162
+ region=region,
163
+ distance_type=distance_type,
164
+ auth_credentials=auth_credentials,
165
+ bigquery_client=bigquery_client,
166
+ )
167
+
168
+ @property
169
+ def client(self) -> Union[bigquery.Client, None]:
170
+ """Return the BigQuery client."""
171
+ if not self._client:
172
+ return None
173
+ return self._client
174
+
175
+ @staticmethod
176
+ def _initialize_client(
177
+ project_id: Union[str, None],
178
+ region: Union[str, None],
179
+ auth_credentials: Union[credentials.Credentials, None],
180
+ ) -> bigquery.Client:
181
+ """
182
+ Initialize a new BigQuery client using the provided `project_id`, `region` and/or `auth_credentials`.
183
+ Defaults will be used in place of missing arguments. For details on authentication, see:
184
+ https://googleapis.dev/python/google-api-core/latest/auth.html
185
+
186
+ Args:
187
+ project_id: GCP project ID for the new client, or None to use default project resolution.
188
+ region: GCP region for the new client, or None to use default region.
189
+ auth_credentials: Credentials to authenticate the new client, or None to use default credentials.
190
+
191
+ Returns:
192
+ An initialized BigQuery client.
193
+
194
+ """
195
+ return bigquery.Client(
196
+ project=project_id or None,
197
+ location=region or None,
198
+ credentials=auth_credentials or None,
199
+ )
200
+
201
+ @staticmethod
202
+ def _bigquery_row_to_node(row: _BigQueryRow) -> BaseNode:
203
+ """
204
+ Convert a BigQuery row to a BaseNode object.
205
+
206
+ Args:
207
+ row: A row retrieved from BigQuery containing node_id, text,
208
+ metadata, embedding, and optional distance.
209
+
210
+ Returns:
211
+ Node object.
212
+
213
+ """
214
+ node_id: str = row.node_id
215
+ text: str = row.text
216
+ metadata: Dict[str, Any] = row.metadata
217
+ embedding: List[float] = row.embedding
218
+ _: Union[float, None] = row.distance
219
+
220
+ try:
221
+ node = metadata_dict_to_node(metadata)
222
+ node.set_content(text)
223
+ node.embedding = embedding
224
+ except (ValueError, TypeError) as e:
225
+ node = TextNode(
226
+ id_=node_id,
227
+ text=text,
228
+ metadata=metadata,
229
+ embedding=embedding,
230
+ )
231
+ _logger.warning(
232
+ f"Failed to construct node {node_id} from metadata. Falling back to manual construction. Error: {e}"
233
+ )
234
+
235
+ return node
236
+
237
+ def _create_dataset_if_not_exists(self, dataset_id: str) -> bigquery.Dataset:
238
+ """
239
+ Create a BigQuery dataset if it does not already exist.
240
+
241
+ For more details on creating datasets, visit:
242
+ https://cloud.google.com/bigquery/docs/datasets#create-dataset
243
+
244
+ Args:
245
+ dataset_id: The ID of the dataset to create.
246
+
247
+ Returns:
248
+ Dataset ID.
249
+
250
+ """
251
+ dataset_ref = bigquery.dataset.DatasetReference(
252
+ project=self._client.project, dataset_id=dataset_id
253
+ )
254
+
255
+ return self._client.create_dataset(dataset_ref, exists_ok=True)
256
+
257
+ def _create_table_if_not_exists(self, table_id) -> bigquery.Table:
258
+ """
259
+ Create a BigQuery table if it does not already exist.
260
+
261
+ For more information on creating tables, visit:
262
+ https://cloud.google.com/bigquery/docs/tables#create-table
263
+
264
+ Args:
265
+ table_id: The ID of the table to create.
266
+
267
+ Returns:
268
+ BigQuery table instance.
269
+
270
+ """
271
+ schema = [
272
+ bigquery.SchemaField("node_id", "STRING", mode="REQUIRED"),
273
+ bigquery.SchemaField("text", "STRING", mode="REQUIRED"),
274
+ bigquery.SchemaField("metadata", "JSON"),
275
+ bigquery.SchemaField("embedding", "FLOAT", mode="REPEATED"),
276
+ ]
277
+
278
+ table_ref = bigquery.TableReference.from_string(
279
+ f"{self._client.project}.{self._dataset.dataset_id}.{table_id}"
280
+ )
281
+ to_create = bigquery.Table(table_ref, schema=schema)
282
+
283
+ return self._client.create_table(to_create, exists_ok=True)
284
+
285
+ def add(self, nodes: List[BaseNode], **add_kwargs: Any) -> List[str]:
286
+ """
287
+ Add nodes to index.
288
+
289
+ Args:
290
+ nodes: List of nodes with embeddings.
291
+
292
+ Returns:
293
+ List of node IDs that were added.
294
+
295
+ """
296
+ node_ids: List[str] = []
297
+ json_records: List[Dict[str, Any]] = []
298
+
299
+ for node in nodes:
300
+ record = {
301
+ "node_id": node.node_id,
302
+ "text": node.get_content(metadata_mode=MetadataMode.NONE),
303
+ "embedding": node.get_embedding(),
304
+ "metadata": node_to_metadata_dict(
305
+ node, remove_text=True, flat_metadata=False
306
+ ),
307
+ }
308
+ node_ids.append(node.node_id)
309
+ json_records.append(record)
310
+
311
+ job_config = bigquery.LoadJobConfig(schema=self._table.schema)
312
+ job = self._client.load_table_from_json(
313
+ json_rows=json_records, destination=self._table, job_config=job_config
314
+ )
315
+ job.result()
316
+
317
+ return node_ids
318
+
319
+ def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
320
+ """
321
+ Delete nodes using with ref_doc_id.
322
+
323
+ Args:
324
+ ref_doc_id : The doc_id of the document to delete.
325
+
326
+ """
327
+ query = f"""
328
+ DELETE FROM `{self._full_table_id}`
329
+ WHERE SAFE.JSON_VALUE(metadata, '$."doc_id"') = @to_delete;
330
+ """
331
+ job_config = bigquery.QueryJobConfig(
332
+ query_parameters=[
333
+ bigquery.ScalarQueryParameter(
334
+ name="to_delete", type_="STRING", value=ref_doc_id
335
+ ),
336
+ ]
337
+ )
338
+
339
+ self._client.query_and_wait(query, job_config=job_config)
340
+
341
+ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:
342
+ """
343
+ Query the vector store using BigQuery's VECTOR_SEARCH to retrieve the top-k most similar nodes.
344
+
345
+ When `MetadataFilters` are provided and the table is indexed on relevant columns, BigQuery attempts to optimize
346
+ the search with pre-filtering before nearest neighbor search. If filters don't align with an index,
347
+ post-filtering is applied after similarity search, potentially returning fewer than `similarity_top_k results`.
348
+ Consider increasing `similarity_top_k` when post-filtering is expected.
349
+
350
+ For more information on pre-filtering and post-filtering, see:
351
+ https://cloud.google.com/bigquery/docs/vector-index#pre-filters_and_post-filters
352
+
353
+ Assumes embeddings are normalized for similarity scoring.
354
+
355
+ Args:
356
+ query: Contains the query embedding, similarity_top_k value, and optional metadata filters.
357
+
358
+ Returns:
359
+ VectorStoreQueryResult
360
+
361
+ """
362
+ where_clause, query_params = build_where_clause_and_params(
363
+ filters=query.filters, node_ids=query.node_ids
364
+ )
365
+
366
+ base_table_query = f"""
367
+ SELECT
368
+ node_id,
369
+ text,
370
+ metadata,
371
+ embedding
372
+ FROM `{self._full_table_id}`
373
+ """
374
+
375
+ if where_clause:
376
+ base_table_query += f" WHERE {where_clause}"
377
+
378
+ query_table_query = f"SELECT {query.query_embedding} AS input_embedding"
379
+
380
+ vector_search_query = f"""
381
+ SELECT base.node_id AS node_id,
382
+ base.text AS text,
383
+ base.metadata AS metadata,
384
+ base.embedding AS embedding,
385
+ distance
386
+ FROM
387
+ VECTOR_SEARCH(
388
+ ({base_table_query}), 'embedding',
389
+ ({query_table_query}), 'input_embedding',
390
+ top_k => @top_k,
391
+ distance_type => @distance_type
392
+ );
393
+ """
394
+
395
+ query_params.extend(
396
+ [
397
+ bigquery.ScalarQueryParameter(
398
+ "top_k", type_="INTEGER", value=query.similarity_top_k
399
+ ),
400
+ bigquery.ScalarQueryParameter(
401
+ "distance_type", type_="STRING", value=self.distance_type
402
+ ),
403
+ ]
404
+ )
405
+ job_config = bigquery.QueryJobConfig(
406
+ query_parameters=query_params,
407
+ )
408
+ rows: bigquery.table.RowIterator = self._client.query_and_wait(
409
+ vector_search_query, job_config=job_config
410
+ )
411
+
412
+ top_k_nodes: List[BaseNode] = []
413
+ top_k_scores: List[float] = []
414
+ top_k_ids: List[str] = []
415
+
416
+ for record in rows:
417
+ row = _BigQueryRow(
418
+ node_id=record.node_id,
419
+ text=record.text,
420
+ metadata=record.metadata,
421
+ embedding=record.embedding,
422
+ distance=record.distance,
423
+ )
424
+ node = self._bigquery_row_to_node(row)
425
+ node_id = record.node_id
426
+ # Assumes embeddings are normalized.
427
+ score = (
428
+ 1 / (1 + record.distance)
429
+ if self.distance_type == DistanceType.EUCLIDEAN
430
+ else (1 + record.distance) / 2
431
+ )
432
+
433
+ top_k_nodes.append(node)
434
+ top_k_scores.append(score)
435
+ top_k_ids.append(node_id)
436
+
437
+ return VectorStoreQueryResult(
438
+ nodes=top_k_nodes, similarities=top_k_scores, ids=top_k_ids
439
+ )
440
+
441
+ def get_nodes(
442
+ self,
443
+ node_ids: Optional[List[str]] = None,
444
+ filters: Optional[MetadataFilters] = None,
445
+ ) -> List[BaseNode]:
446
+ """
447
+ Retrieve nodes from BigQuery using node IDs, metadata filters, or both.
448
+
449
+ If both `node_ids` and `filters` are provided, only nodes that satisfy
450
+ both conditions will be returned.
451
+
452
+ Args:
453
+ node_ids: Optional list of node IDs for retrieval.
454
+ filters : Optional MetadataFilters filters for retrieval.
455
+
456
+ Returns:
457
+ A list of matching nodes.
458
+
459
+ Raises:
460
+ ValueError: If neither `node_ids` nor `filters` is provided.
461
+
462
+ """
463
+ if not (node_ids or filters):
464
+ raise ValueError(
465
+ "get_nodes requires at least one filtering parameter: "
466
+ "'node_ids', 'filters', or both. Received neither."
467
+ )
468
+
469
+ where_clause, query_params = build_where_clause_and_params(node_ids, filters)
470
+
471
+ query = f"""
472
+ SELECT node_id,
473
+ text,
474
+ embedding,
475
+ metadata
476
+ FROM `{self._full_table_id}`
477
+ WHERE {where_clause};
478
+ """
479
+
480
+ job_config = bigquery.QueryJobConfig(
481
+ query_parameters=query_params,
482
+ )
483
+ rows: bigquery.table.RowIterator = self._client.query_and_wait(
484
+ query, job_config=job_config
485
+ )
486
+
487
+ nodes: List[BaseNode] = []
488
+ for record in rows:
489
+ row = _BigQueryRow(
490
+ node_id=record.node_id,
491
+ text=record.text,
492
+ metadata=record.metadata,
493
+ embedding=record.embedding,
494
+ distance=record.distance,
495
+ )
496
+ node = self._bigquery_row_to_node(row)
497
+
498
+ nodes.append(node)
499
+
500
+ return nodes
501
+
502
+ def delete_nodes(
503
+ self,
504
+ node_ids: Optional[List[str]] = None,
505
+ filters: Optional[MetadataFilters] = None,
506
+ **delete_kwargs: Any,
507
+ ) -> None:
508
+ """
509
+ Delete nodes from BigQuery based on node IDs, metadata filters, or both.
510
+
511
+ If both `node_ids` and `filters` are provided, only nodes matching both
512
+ criteria will be deleted.
513
+
514
+ Args:
515
+ node_ids: Optional list of node IDs to delete.
516
+ filters : Optional MetadataFilters filters for deletion.
517
+
518
+ Raises:
519
+ ValueError: If neither `node_ids` nor `filters` are provided.
520
+
521
+ """
522
+ if not (node_ids or filters):
523
+ raise ValueError(
524
+ "delete_nodes requires at least one filtering parameter: "
525
+ "'node_ids', 'filters', or both. Received neither."
526
+ )
527
+
528
+ where_clause, query_params = build_where_clause_and_params(node_ids, filters)
529
+
530
+ query = f"""
531
+ DELETE FROM `{self._full_table_id}`
532
+ WHERE {where_clause};
533
+ """
534
+
535
+ job_config = bigquery.QueryJobConfig(
536
+ query_parameters=query_params,
537
+ )
538
+ self._client.query_and_wait(query, job_config=job_config)
539
+
540
+ def clear(self) -> None:
541
+ """
542
+ Clears the index.
543
+
544
+ This truncates the underlying table in BigQuery.
545
+ """
546
+ query = f"""TRUNCATE TABLE `{self._full_table_id}`;"""
547
+ self._client.query_and_wait(query)
@@ -0,0 +1,162 @@
1
+ from typing import List, Optional, Tuple, Union
2
+
3
+ from google.cloud import bigquery
4
+ from llama_index.core.vector_stores.types import (
5
+ FilterOperator,
6
+ MetadataFilter,
7
+ MetadataFilters,
8
+ )
9
+
10
+
11
+ def build_where_clause_and_params(
12
+ node_ids: Optional[List[str]] = None,
13
+ filters: Optional[MetadataFilters] = None,
14
+ ) -> Tuple[
15
+ str, List[Union[bigquery.ScalarQueryParameter, bigquery.ArrayQueryParameter]]
16
+ ]:
17
+ """
18
+ Construct a parameterized SQL WHERE clause and corresponding query parameters.
19
+
20
+ The clause is built from the provided node IDs and metadata filters. Parameters
21
+ are returned separately to support safe, parameterized queries in BigQuery,
22
+ helping to prevent SQL injection. See:
23
+ https://cloud.google.com/bigquery/docs/parameterized-queries
24
+
25
+ If both `node_ids` and `filters` are provided, the resulting WHERE clause
26
+ combines conditions using AND logic.
27
+
28
+ Args:
29
+ node_ids: A list of node IDs to include in the filter.
30
+ filters: Metadata filters to apply to the query.
31
+
32
+ Returns:
33
+ A tuple (where_clause, query_params), where where_clause is a SQL WHERE clause string,
34
+ and query_params is a list of query parameters to bind to the query.
35
+
36
+ """
37
+ query_params: List[
38
+ Union[bigquery.ScalarQueryParameter, bigquery.ArrayQueryParameter]
39
+ ] = []
40
+ conditions: List[str] = []
41
+
42
+ if filters:
43
+ filter_where_clause, filter_query_params = (
44
+ _recursive_build_where_clause_from_filters(filters)
45
+ )
46
+ conditions.append(filter_where_clause)
47
+ query_params.extend(filter_query_params)
48
+
49
+ if node_ids:
50
+ conditions.append("node_id IN UNNEST(@node_ids)")
51
+ query_params.append(
52
+ bigquery.ArrayQueryParameter(
53
+ name="node_ids", array_type="STRING", values=node_ids
54
+ )
55
+ )
56
+
57
+ # if both `node_ids` and `filters` are provided, both criteria should be considered
58
+ where_clause = " AND ".join(conditions)
59
+
60
+ return where_clause, query_params
61
+
62
+
63
+ def _recursive_build_where_clause_from_filters(
64
+ meta_filters: MetadataFilters,
65
+ ) -> Tuple[
66
+ str, List[Union[bigquery.ScalarQueryParameter, bigquery.ArrayQueryParameter]]
67
+ ]:
68
+ """
69
+ Recursively construct a SQL WHERE clause and corresponding query parameters
70
+
71
+ The provided MetadataFilters object may contain nested filter groups. This function
72
+ traverses them recursively to build a complete WHERE clause and the associated
73
+ query parameters for use in a parameterized BigQuery query.
74
+
75
+ Args:
76
+ meta_filters: A potentially nested MetadataFilters filter object
77
+
78
+ Returns:
79
+ A tuple (where_clause, query_params), where where_clause is a parameterized SQL WHERE clause string,
80
+ and query_params is a list of query parameters to bind to the query.
81
+
82
+ """
83
+ filters_list: List[str] = []
84
+ query_params: List[
85
+ Union[bigquery.ScalarQueryParameter, bigquery.ArrayQueryParameter]
86
+ ] = []
87
+
88
+ for filter_ in meta_filters.filters:
89
+ clause, params = (
90
+ _recursive_build_where_clause_from_filters(filter_)
91
+ if isinstance(filter_, MetadataFilters)
92
+ else _build_filter_clause(filter_)
93
+ )
94
+
95
+ filters_list.append(clause)
96
+ query_params.extend(params)
97
+
98
+ condition = f" {meta_filters.condition.value.upper()} "
99
+ filters_ = f"({condition.join(filters_list)})"
100
+
101
+ return filters_, query_params
102
+
103
+
104
+ def _build_filter_clause(
105
+ filter_: MetadataFilter,
106
+ ) -> Tuple[
107
+ str, List[Union[bigquery.ScalarQueryParameter, bigquery.ArrayQueryParameter]]
108
+ ]:
109
+ field = filter_.key
110
+ operator = filter_.operator
111
+ value = filter_.value
112
+
113
+ if operator == FilterOperator.IS_EMPTY:
114
+ clause = f"JSON_TYPE(JSON_QUERY(metadata, '$.\"{field}\"')) = 'null'"
115
+ params = []
116
+ elif operator == FilterOperator.IN or operator == FilterOperator.NIN:
117
+ bigquery_operator = _llama_to_bigquery_operator(operator)
118
+ clause = (
119
+ f"SAFE.JSON_VALUE(metadata, '$.\"{field}\"') {bigquery_operator} UNNEST(?)"
120
+ )
121
+ params = [
122
+ bigquery.ArrayQueryParameter(name=None, array_type="STRING", values=value)
123
+ ]
124
+ elif operator == FilterOperator.TEXT_MATCH:
125
+ bigquery_operator = _llama_to_bigquery_operator(operator)
126
+ clause = (
127
+ f"SAFE.JSON_VALUE(metadata, '$.\"{field}\"') {bigquery_operator} '{value}'"
128
+ )
129
+ params = [bigquery.ScalarQueryParameter(name=None, type_="STRING", value=value)]
130
+ elif operator == FilterOperator.TEXT_MATCH_INSENSITIVE:
131
+ bigquery_operator = _llama_to_bigquery_operator(operator)
132
+ clause = f"LOWER(SAFE.JSON_VALUE(metadata, '$.\"{field}\"')) {bigquery_operator} LOWER('{value}')"
133
+ params = [bigquery.ScalarQueryParameter(name=None, type_="STRING", value=value)]
134
+ else:
135
+ bigquery_operator = _llama_to_bigquery_operator(operator)
136
+ clause = f"SAFE.JSON_VALUE(metadata, '$.\"{field}\"') {bigquery_operator} ?"
137
+ params = [bigquery.ScalarQueryParameter(name=None, type_="STRING", value=value)]
138
+
139
+ return clause, params
140
+
141
+
142
+ def _llama_to_bigquery_operator(operator: FilterOperator) -> str:
143
+ operator_map = {
144
+ FilterOperator.EQ: "=",
145
+ FilterOperator.GT: ">",
146
+ FilterOperator.LT: "<",
147
+ FilterOperator.NE: "!=",
148
+ FilterOperator.GTE: ">=",
149
+ FilterOperator.LTE: "<=",
150
+ FilterOperator.IN: "IN",
151
+ FilterOperator.NIN: "NOT IN",
152
+ FilterOperator.TEXT_MATCH: "LIKE",
153
+ FilterOperator.TEXT_MATCH_INSENSITIVE: "LIKE",
154
+ FilterOperator.IS_EMPTY: "IS NULL",
155
+ }
156
+
157
+ try:
158
+ return operator_map[operator]
159
+ except KeyError:
160
+ raise ValueError(
161
+ f"Invalid operator `{operator.value}` is not a supported BigQuery operator."
162
+ )
@@ -0,0 +1,64 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [dependency-groups]
6
+ dev = [
7
+ "ipython==8.10.0",
8
+ "jupyter>=1.0.0,<2",
9
+ "mypy==0.991",
10
+ "pre-commit==3.2.0",
11
+ "pylint==2.15.10",
12
+ "pytest==7.2.1",
13
+ "pytest-mock==3.11.1",
14
+ "ruff==0.11.11",
15
+ "types-Deprecated>=0.1.0",
16
+ "types-PyYAML>=6.0.12.12,<7",
17
+ "types-protobuf>=4.24.0.4,<5",
18
+ "types-redis==4.5.5.0",
19
+ "types-requests==2.28.11.8",
20
+ "types-setuptools==67.1.0.0",
21
+ "black[jupyter]<=23.9.1,>=23.7.0",
22
+ "codespell[toml]>=v2.2.6",
23
+ "diff-cover>=9.2.0",
24
+ "pytest-cov>=6.1.1",
25
+ "sqlparse>=0.5.3",
26
+ ]
27
+
28
+ [project]
29
+ name = "llama-index-vector-stores-bigquery"
30
+ version = "0.1.0"
31
+ description = "llama-index vector_stores bigquery vector search integration"
32
+ requires-python = ">=3.9,<4.0"
33
+ readme = "README.md"
34
+ license = "MIT"
35
+ dependencies = [
36
+ "google-cloud-bigquery>=3.34.0,<4",
37
+ "llama-index-core>=0.13.0,<0.14",
38
+ ]
39
+
40
+ [tool.codespell]
41
+ check-filenames = true
42
+ check-hidden = true
43
+ skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
44
+
45
+ [tool.hatch.build.targets.sdist]
46
+ include = ["llama_index/"]
47
+ exclude = ["**/BUILD"]
48
+
49
+ [tool.hatch.build.targets.wheel]
50
+ include = ["llama_index/"]
51
+ exclude = ["**/BUILD"]
52
+
53
+ [tool.llamahub]
54
+ contains_example = false
55
+ import_path = "llama_index.vector_stores.bigquery"
56
+
57
+ [tool.llamahub.class_authors]
58
+ BigQueryVectorStore = "llama-index"
59
+
60
+ [tool.mypy]
61
+ disallow_untyped_defs = true
62
+ exclude = ["_static", "build", "examples", "notebooks", "venv"]
63
+ ignore_missing_imports = true
64
+ python_version = "3.8"