From 5db79c208afb38de2e42ec0360d5d7303a913500 Mon Sep 17 00:00:00 2001 From: Siddhesh Khairnar Date: Tue, 7 Jul 2026 11:34:06 +0530 Subject: [PATCH] feat: Add Pinecone online store integration Signed-off-by: Siddhesh Khairnar --- docs/SUMMARY.md | 1 + docs/reference/alpha-vector-database.md | 9 +- docs/reference/online-stores/README.md | 4 + docs/reference/online-stores/pinecone.md | 109 +++ pixi.lock | 7 +- pyproject.toml | 3 +- .../pinecone_online_store/__init__.py | 6 + .../pinecone_online_store/pinecone.py | 625 ++++++++++++++++++ .../pinecone_repo_configuration.py | 12 + sdk/python/feast/repo_config.py | 1 + .../unit/infra/online_store/test_pinecone.py | 518 +++++++++++++++ .../universal/online_store/pinecone.py | 27 + 12 files changed, 1317 insertions(+), 5 deletions(-) create mode 100644 docs/reference/online-stores/pinecone.md create mode 100644 sdk/python/feast/infra/online_stores/pinecone_online_store/__init__.py create mode 100644 sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone.py create mode 100644 sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone_repo_configuration.py create mode 100644 sdk/python/tests/unit/infra/online_store/test_pinecone.py create mode 100644 sdk/python/tests/universal/feature_repos/universal/online_store/pinecone.py diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md index d35dbf1652b..477c8e57c63 100644 --- a/docs/SUMMARY.md +++ b/docs/SUMMARY.md @@ -166,6 +166,7 @@ * [Aerospike](reference/online-stores/aerospike.md) * [Elasticsearch](reference/online-stores/elasticsearch.md) * [Qdrant](reference/online-stores/qdrant.md) + * [Pinecone](reference/online-stores/pinecone.md) * [Faiss](reference/online-stores/faiss.md) * [Hybrid](reference/online-stores/hybrid.md) * [Registries](reference/registries/README.md) diff --git a/docs/reference/alpha-vector-database.md b/docs/reference/alpha-vector-database.md index 28d0bf0098e..68613e8b02e 100644 --- a/docs/reference/alpha-vector-database.md +++ b/docs/reference/alpha-vector-database.md @@ -16,6 +16,7 @@ Below are supported vector databases and implemented features: | SQLite | [x] | [ ] | [x] | [x] | | Qdrant | [x] | [x] | [] | [] | | ScyllaDB | [x] | [x] | [x] | [x] | +| Pinecone | [x] | [x] | [x] | [x] | *Note: V2 Support means the SDK supports retrieval of features along with vector embeddings from vector similarity search. @@ -190,7 +191,7 @@ print('\n'.join([c.message.content for c in response.choices])) ### Configuration and Installation -We offer [Milvus](https://milvus.io/), [PGVector](https://github.com/pgvector/pgvector), [SQLite](https://github.com/asg017/sqlite-vec), [Elasticsearch](https://www.elastic.co) and [Qdrant](https://qdrant.tech/) as Online Store options for Vector Databases. +We offer [Milvus](https://milvus.io/), [PGVector](https://github.com/pgvector/pgvector), [SQLite](https://github.com/asg017/sqlite-vec), [Elasticsearch](https://www.elastic.co), [Qdrant](https://qdrant.tech/) and [Pinecone](https://www.pinecone.io/) as Online Store options for Vector Databases. Milvus offers a convenient local implementation for vector similarity search. To use Milvus, you can install the Feast package with the Milvus extra. @@ -210,6 +211,12 @@ pip install feast[elasticsearch] ```bash pip install feast[qdrant] ``` +#### Installation with Pinecone + +```bash +pip install feast[pinecone] +``` + #### Installation with SQLite If you are using `pyenv` to manage your Python versions, you can install the SQLite extension with the following command: diff --git a/docs/reference/online-stores/README.md b/docs/reference/online-stores/README.md index 39294966170..d0f80302701 100644 --- a/docs/reference/online-stores/README.md +++ b/docs/reference/online-stores/README.md @@ -90,6 +90,10 @@ Please see [Online Store](../../getting-started/components/online-store.md) for [milvus.md](milvus.md) {% endcontent-ref %} +{% content-ref url="pinecone.md" %} +[pinecone.md](pinecone.md) +{% endcontent-ref %} + {% content-ref url="faiss.md" %} [faiss.md](faiss.md) {% endcontent-ref %} diff --git a/docs/reference/online-stores/pinecone.md b/docs/reference/online-stores/pinecone.md new file mode 100644 index 00000000000..2c03b0a7a59 --- /dev/null +++ b/docs/reference/online-stores/pinecone.md @@ -0,0 +1,109 @@ +# Pinecone online store + +## Description + +The [Pinecone](https://www.pinecone.io/) online store provides support for materializing feature values into a Pinecone vector database. It is particularly suited for AI workloads that require both low-latency feature serving and vector similarity search for RAG (Retrieval-Augmented Generation) applications. + +## Getting started + +In order to use this online store, you'll need to install the Pinecone extra (along with the dependency needed for the offline store of choice). E.g. + +`pip install 'feast[pinecone]'` + +You will also need a Pinecone account and API key. Set your API key via the `PINECONE_API_KEY` environment variable or directly in `feature_store.yaml`. + +## Examples + +{% code title="feature_store.yaml" %} +```yaml +project: my_feature_repo +registry: data/registry.db +provider: local +online_store: + type: pinecone + api_key: ${PINECONE_API_KEY} + index_name: feast-online + embedding_dim: 384 + metric: cosine + cloud: aws + region: us-east-1 + vector_enabled: true +``` +{% endcode %} + +### Using a custom namespace + +By default, Feast maps each feature view to a Pinecone namespace (`{project}_{feature_view_name}`). You can override this with a fixed namespace: + +{% code title="feature_store.yaml" %} +```yaml +online_store: + type: pinecone + api_key: ${PINECONE_API_KEY} + index_name: feast-online + namespace: my-custom-namespace + embedding_dim: 384 +``` +{% endcode %} + +### Vector similarity search + +Pinecone supports the `retrieve_online_documents_v2` API for vector similarity search: + +```python +from feast import FeatureStore + +store = FeatureStore(repo_path=".") + +results = store.retrieve_online_documents_v2( + features=[ + "city_embeddings:vector", + "city_embeddings:sentence_chunks", + ], + query=query_embedding, + top_k=5, + distance_metric="cosine", +).to_df() +``` + +The full set of configuration options is available in [PineconeOnlineStoreConfig](https://rtd.feast.dev/en/latest/#feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStoreConfig). + +## Configuration options + +| Option | Type | Default | Description | +|:---|:---|:---|:---| +| `type` | string | `"pinecone"` | Must be `pinecone` | +| `api_key` | string | `None` | Pinecone API key (falls back to `PINECONE_API_KEY` env var) | +| `index_name` | string | `"feast-online"` | Name of the Pinecone index | +| `namespace` | string | `None` | Override namespace (default: `{project}_{fv_name}`) | +| `embedding_dim` | int | `128` | Vector dimension | +| `metric` | string | `"cosine"` | Distance metric: `cosine`, `euclidean`, or `dotproduct` | +| `cloud` | string | `"aws"` | Cloud provider for serverless indexes | +| `region` | string | `"us-east-1"` | Cloud region for serverless indexes | +| `vector_enabled` | bool | `true` | Enable vector similarity search | +| `batch_size` | int | `100` | Number of vectors per upsert batch | + +## Functionality Matrix + +The set of functionality supported by online stores is described in detail [here](overview.md#functionality). +Below is a matrix indicating which functionality is supported by the Pinecone online store. + +| | Pinecone | +|:----------------------------------------------------------|:---------| +| write feature values to the online store | yes | +| read feature values from the online store | yes | +| update infrastructure (e.g. tables) in the online store | yes | +| teardown infrastructure (e.g. tables) in the online store | yes | +| generate a plan of infrastructure changes | no | +| support for on-demand transforms | yes | +| readable by Python SDK | yes | +| readable by Java | no | +| readable by Go | no | +| support for entityless feature views | yes | +| support for concurrent writing to the same key | yes | +| support for ttl (time to live) at retrieval | no | +| support for deleting expired data | no | +| collocated by feature view | no | +| collocated by feature service | no | +| collocated by entity key | no | +| vector similarity search | yes | diff --git a/pixi.lock b/pixi.lock index b9f1135da9c..87c813eb958 100644 --- a/pixi.lock +++ b/pixi.lock @@ -2424,8 +2424,8 @@ packages: requires_python: '>=3.10' - pypi: ./ name: feast - version: 0.64.1.dev40+g82a6918c1.d20260717 - sha256: 5ab6cffee504478fb0d04132e90e2ea5afa7776ff5e0cd695853e25e71239657 + version: 0.1.dev4624+gf487b37fd.d20260718 + sha256: 9b12db0059ce0dc18ba8791e31783a37944bd10df17bca1cf7a6417fe85f8376 requires_dist: - attrs - click>=7.0.0,<9.0.0 @@ -2518,6 +2518,7 @@ packages: - psycopg[c,pool]>=3.2.5 ; extra == 'postgres-c' - torch>=2.7.0 ; extra == 'pytorch' - torchvision>=0.22.1 ; extra == 'pytorch' + - pinecone>=5.0 ; extra == 'pinecone' - qdrant-client>=1.12.0 ; extra == 'qdrant' - transformers>=4.36.0 ; extra == 'rag' - datasets>=3.6.0 ; extra == 'rag' @@ -2549,7 +2550,7 @@ packages: - minio==7.2.11 ; extra == 'test' - python-keycloak==4.2.2 ; extra == 'test' - cryptography>=43.0 ; extra == 'test' - - feast[aws,azure,cassandra,clickhouse,couchbase,delta,docling,duckdb,elasticsearch,faiss,gcp,ge,go,grpcio,hazelcast,hbase,ibis,iceberg,image,k8s,mcp,milvus,mlflow,mongodb,mssql,mysql,openlineage,opentelemetry,oracle,postgres,pytorch,qdrant,rag,ray,redis,scylladb,singlestore,snowflake,spark,sqlite-vec,test,trino] ; extra == 'ci' + - feast[aws,azure,cassandra,clickhouse,couchbase,delta,docling,duckdb,elasticsearch,faiss,gcp,ge,go,grpcio,hazelcast,hbase,ibis,iceberg,image,k8s,mcp,milvus,mlflow,mongodb,mssql,mysql,openlineage,opentelemetry,oracle,pinecone,postgres,pytorch,qdrant,rag,ray,redis,scylladb,singlestore,snowflake,spark,sqlite-vec,test,trino] ; extra == 'ci' - build ; extra == 'ci' - virtualenv==20.23.0 ; extra == 'ci' - dbt-artifacts-parser ; extra == 'ci' diff --git a/pyproject.toml b/pyproject.toml index 6675ef16215..a6edc6d9331 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,6 +119,7 @@ postgres = ["psycopg[binary,pool]>=3.2.5"] # https://www.psycopg.org/psycopg3/docs/basic/install.html#local-installation postgres-c = ["psycopg[c,pool]>=3.2.5"] pytorch = ["torch>=2.7.0", "torchvision>=0.22.1"] +pinecone = ["pinecone>=5.0"] qdrant = ["qdrant-client>=1.12.0"] rag = [ "transformers>=4.36.0", @@ -164,7 +165,7 @@ test = [ ] ci = [ - "feast[test, aws, azure, cassandra, clickhouse, couchbase, delta, docling, duckdb, elasticsearch, faiss, gcp, ge, go, grpcio, hazelcast, hbase, ibis, iceberg, image, k8s, mcp, milvus, mlflow, mongodb, mssql, mysql, openlineage, opentelemetry, oracle, scylladb, spark, trino, postgres, pytorch, qdrant, rag, ray, redis, singlestore, snowflake, sqlite_vec]", + "feast[test, aws, azure, cassandra, clickhouse, couchbase, delta, docling, duckdb, elasticsearch, faiss, gcp, ge, go, grpcio, hazelcast, hbase, ibis, iceberg, image, k8s, mcp, milvus, mlflow, mongodb, mssql, mysql, openlineage, opentelemetry, oracle, pinecone, scylladb, spark, trino, postgres, pytorch, qdrant, rag, ray, redis, singlestore, snowflake, sqlite_vec]", "build", "virtualenv==20.23.0", "dbt-artifacts-parser", diff --git a/sdk/python/feast/infra/online_stores/pinecone_online_store/__init__.py b/sdk/python/feast/infra/online_stores/pinecone_online_store/__init__.py new file mode 100644 index 00000000000..bfdda9e8db5 --- /dev/null +++ b/sdk/python/feast/infra/online_stores/pinecone_online_store/__init__.py @@ -0,0 +1,6 @@ +from feast.infra.online_stores.pinecone_online_store.pinecone import ( + PineconeOnlineStore, + PineconeOnlineStoreConfig, +) + +__all__ = ["PineconeOnlineStore", "PineconeOnlineStoreConfig"] diff --git a/sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone.py b/sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone.py new file mode 100644 index 00000000000..e8d0ddb935b --- /dev/null +++ b/sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone.py @@ -0,0 +1,625 @@ +import base64 +import logging +import os +import time +from datetime import datetime +from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple + +from pydantic import StrictStr + +from feast import Entity +from feast.feature_view import FeatureView +from feast.infra.infra_object import InfraObject +from feast.infra.key_encoding_utils import ( + deserialize_entity_key, + serialize_entity_key, +) +from feast.infra.online_stores.helpers import compute_table_id +from feast.infra.online_stores.online_store import OnlineStore +from feast.infra.online_stores.vector_store import VectorStoreConfig +from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.repo_config import FeastConfigBaseModel, RepoConfig +from feast.utils import ( + _serialize_vector_to_float_list, + to_naive_utc, +) + +logger = logging.getLogger(__name__) + +PINECONE_METRIC_MAPPING = { + "cosine": "cosine", + "COSINE": "cosine", + "l2": "euclidean", + "L2": "euclidean", + "euclidean": "euclidean", + "EUCLIDEAN": "euclidean", + "dot": "dotproduct", + "DOT": "dotproduct", + "dotproduct": "dotproduct", + "DOTPRODUCT": "dotproduct", + "IP": "dotproduct", +} + +# Max metadata size per vector in Pinecone (40 KB) +_MAX_METADATA_BYTES = 40_960 + +# Pinecone upsert batch size limit +_UPSERT_BATCH_SIZE = 100 + +# Max wait time for index readiness (seconds) +_INDEX_READY_TIMEOUT = 300 + + +class PineconeOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig): + """ + Configuration for the Pinecone online store. + + NOTE: The class *must* end with the ``OnlineStoreConfig`` suffix. + """ + + type: Literal["pinecone"] = "pinecone" + + api_key: Optional[StrictStr] = None + """Pinecone API key. Falls back to the ``PINECONE_API_KEY`` env var.""" + + index_name: Optional[StrictStr] = "feast-online" + """Name of the Pinecone index to use.""" + + namespace: Optional[StrictStr] = None + """Default Pinecone namespace. When ``None``, the namespace is derived + from the project name and feature view (``{project}_{fv_name}``).""" + + embedding_dim: Optional[int] = 128 + """Dimensionality of vectors stored in the index.""" + + metric: Optional[StrictStr] = "cosine" + """Distance metric: ``cosine``, ``euclidean``, or ``dotproduct``.""" + + cloud: Optional[StrictStr] = "aws" + """Cloud provider for serverless indexes (``aws``, ``gcp``, ``azure``).""" + + region: Optional[StrictStr] = "us-east-1" + """Cloud region for serverless indexes.""" + + vector_enabled: Optional[bool] = True + """Whether vector similarity search is enabled.""" + + batch_size: Optional[int] = _UPSERT_BATCH_SIZE + """Number of vectors per upsert request.""" + + +class PineconeOnlineStore(OnlineStore): + """ + Pinecone implementation of the Feast online store interface. + + Stores feature values as Pinecone vectors with metadata. Each feature + view is mapped to a Pinecone namespace within a single index. + """ + + _client: Optional[Any] = None + _index: Optional[Any] = None + + def _get_api_key(self, config: RepoConfig) -> str: + online_config = config.online_store + assert isinstance(online_config, PineconeOnlineStoreConfig) + api_key = online_config.api_key or os.environ.get("PINECONE_API_KEY") + if not api_key: + raise ValueError( + "Pinecone API key is required. Set it in feature_store.yaml " + "(online_store.api_key) or via the PINECONE_API_KEY env var." + ) + return api_key + + def _get_client(self, config: RepoConfig) -> Any: + if self._client is not None: + return self._client + + from pinecone import Pinecone + + api_key = self._get_api_key(config) + self._client = Pinecone(api_key=api_key) + return self._client + + def _get_index(self, config: RepoConfig) -> Any: + if self._index is not None: + return self._index + + online_config = config.online_store + assert isinstance(online_config, PineconeOnlineStoreConfig) + client = self._get_client(config) + self._index = client.Index(online_config.index_name) + return self._index + + def _get_namespace(self, config: RepoConfig, table: FeatureView) -> str: + online_config = config.online_store + assert isinstance(online_config, PineconeOnlineStoreConfig) + if online_config.namespace: + return online_config.namespace + return _table_id( + config.project, + table, + config.registry.enable_online_feature_view_versioning, + ) + + def online_write_batch( + self, + config: RepoConfig, + table: FeatureView, + data: List[ + Tuple[ + EntityKeyProto, + Dict[str, ValueProto], + datetime, + Optional[datetime], + ] + ], + progress: Optional[Callable[[int], Any]], + ) -> None: + index = self._get_index(config) + namespace = self._get_namespace(config, table) + online_config = config.online_store + assert isinstance(online_config, PineconeOnlineStoreConfig) + + vector_fields = {f.name for f in table.schema if f.vector_index} + + # De-duplicate: keep only the latest event per entity key. + unique_entities: Dict[ + str, + Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]], + ] = {} + for entity_key, values_dict, timestamp, created_ts in data: + entity_key_str = serialize_entity_key( + entity_key, + entity_key_serialization_version=config.entity_key_serialization_version, + ).hex() + + existing = unique_entities.get(entity_key_str) + if existing is None or existing[2] < timestamp: + unique_entities[entity_key_str] = ( + entity_key, + values_dict, + timestamp, + created_ts, + ) + + vectors_to_upsert = [] + for entity_key_str, ( + entity_key, + values_dict, + timestamp, + created_ts, + ) in unique_entities.items(): + timestamp_utc = to_naive_utc(timestamp) + created_ts_utc = to_naive_utc(created_ts) if created_ts else None + + metadata: Dict[str, Any] = { + "event_ts": int(timestamp_utc.timestamp() * 1e6), + "created_ts": int(created_ts_utc.timestamp() * 1e6) + if created_ts_utc + else 0, + "entity_key": entity_key_str, + } + + embedding: Optional[List[float]] = None + + for feature_name, value_proto in values_dict.items(): + if feature_name in vector_fields: + embedding = _extract_vector(value_proto) + else: + metadata[feature_name] = _proto_value_to_metadata(value_proto) + + if embedding is None: + embedding = [0.0] * (online_config.embedding_dim or 128) + + vectors_to_upsert.append( + { + "id": entity_key_str, + "values": embedding, + "metadata": metadata, + } + ) + + if progress: + progress(1) + + batch_size = online_config.batch_size or _UPSERT_BATCH_SIZE + for i in range(0, len(vectors_to_upsert), batch_size): + batch = vectors_to_upsert[i : i + batch_size] + index.upsert(vectors=batch, namespace=namespace) + + def online_read( + self, + config: RepoConfig, + table: FeatureView, + entity_keys: List[EntityKeyProto], + requested_features: Optional[List[str]] = None, + ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: + index = self._get_index(config) + namespace = self._get_namespace(config, table) + + feature_type_map = {f.name: f.dtype for f in table.features} + if getattr(table, "write_to_online_store", False): + feature_type_map.update({f.name: f.dtype for f in table.schema}) + + vector_fields = {f.name for f in table.schema if f.vector_index} + + ids = [] + for entity_key in entity_keys: + entity_key_str = serialize_entity_key( + entity_key, + entity_key_serialization_version=config.entity_key_serialization_version, + ).hex() + ids.append(entity_key_str) + + try: + fetch_response = index.fetch(ids=ids, namespace=namespace) + except Exception: + logger.exception("Error fetching from Pinecone") + return [(None, None)] * len(entity_keys) + + fetched_vectors = fetch_response.get("vectors", {}) + if hasattr(fetch_response, "vectors"): + fetched_vectors = fetch_response.vectors + + result_list: List[ + Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]] + ] = [] + for vector_id in ids: + record = fetched_vectors.get(vector_id) + if record is None: + result_list.append((None, None)) + continue + + metadata = record.get("metadata", {}) + if hasattr(record, "metadata"): + metadata = record.metadata or {} + values = record.get("values", []) + if hasattr(record, "values"): + values = record.values or [] + + event_ts_us = metadata.get("event_ts", 0) + res_ts = datetime.fromtimestamp(event_ts_us / 1e6) if event_ts_us else None + + res: Dict[str, ValueProto] = {} + features_to_read = ( + requested_features + if requested_features + else list(feature_type_map.keys()) + ) + + for feature_name in features_to_read: + if feature_name in vector_fields and values: + val = _serialize_vector_to_float_list(values) + res[feature_name] = val + elif feature_name in metadata: + val = _metadata_to_proto_value( + metadata[feature_name], feature_type_map.get(feature_name) + ) + res[feature_name] = val + + result_list.append((res_ts, res if res else None)) + + return result_list + + def update( + self, + config: RepoConfig, + tables_to_delete: Sequence[FeatureView], + tables_to_keep: Sequence[FeatureView], + entities_to_delete: Sequence[Entity], + entities_to_keep: Sequence[Entity], + partial: bool, + ) -> None: + online_config = config.online_store + assert isinstance(online_config, PineconeOnlineStoreConfig) + client = self._get_client(config) + + index_name = online_config.index_name + existing_indexes = {idx.name for idx in client.list_indexes()} + + if index_name not in existing_indexes: + from pinecone import ServerlessSpec + + metric = PINECONE_METRIC_MAPPING.get( + online_config.metric or "cosine", "cosine" + ) + client.create_index( + name=index_name, + dimension=online_config.embedding_dim or 128, + metric=metric, + spec=ServerlessSpec( + cloud=online_config.cloud or "aws", + region=online_config.region or "us-east-1", + ), + ) + self._wait_for_index_ready(client, index_name or "feast-online") + + # Reset index reference so it re-connects to the (possibly new) index. + self._index = None + + for table in tables_to_delete: + self._delete_namespace(config, table) + + def plan( + self, config: RepoConfig, desired_registry_proto: RegistryProto + ) -> List[InfraObject]: + return [] + + def teardown( + self, + config: RepoConfig, + tables: Sequence[FeatureView], + entities: Sequence[Entity], + ) -> None: + for table in tables: + self._delete_namespace(config, table) + + def retrieve_online_documents_v2( + self, + config: RepoConfig, + table: FeatureView, + requested_features: List[str], + embedding: Optional[List[float]], + top_k: int, + distance_metric: Optional[str] = None, + query_string: Optional[str] = None, + include_feature_view_version_metadata: bool = False, + ) -> List[ + Tuple[ + Optional[datetime], + Optional[EntityKeyProto], + Optional[Dict[str, ValueProto]], + ] + ]: + online_config = config.online_store + assert isinstance(online_config, PineconeOnlineStoreConfig) + + if not online_config.vector_enabled: + raise ValueError("Vector search is not enabled in the online store config") + + if embedding is None and query_string is None: + raise ValueError("Either embedding or query_string must be provided") + + if embedding is None: + raise ValueError( + "Pinecone requires a query embedding for similarity search. " + "Text-only keyword search is not supported." + ) + + index = self._get_index(config) + namespace = self._get_namespace(config, table) + + entity_name_type_map = {k.name: k.dtype for k in table.entity_columns} + vector_fields = {f.name for f in table.schema if f.vector_index} + + pinecone_filter = None + if query_string is not None: + from feast.types import PrimitiveFeastType + from feast.types import ValueType as FeastValueType + + string_fields = [ + f.name + for f in table.features + if isinstance(f.dtype, PrimitiveFeastType) + and f.dtype.to_value_type() == FeastValueType.STRING + and f.name in requested_features + ] + if string_fields: + pinecone_filter = { + "$or": [{field: {"$eq": query_string}} for field in string_fields] + } + + query_response = index.query( + vector=embedding, + top_k=top_k, + namespace=namespace, + include_metadata=True, + include_values=True, + filter=pinecone_filter, + ) + + matches = query_response.get("matches", []) + if hasattr(query_response, "matches"): + matches = query_response.matches or [] + + result_list: List[ + Tuple[ + Optional[datetime], + Optional[EntityKeyProto], + Optional[Dict[str, ValueProto]], + ] + ] = [] + + for match in matches: + metadata = match.get("metadata", {}) + if hasattr(match, "metadata"): + metadata = match.metadata or {} + values = match.get("values", []) + if hasattr(match, "values"): + values = match.values or [] + score = match.get("score", 0.0) + if hasattr(match, "score"): + score = match.score + + event_ts_us = metadata.get("event_ts", 0) + res_ts = datetime.fromtimestamp(event_ts_us / 1e6) if event_ts_us else None + + entity_key_hex = metadata.get("entity_key") + entity_key_proto = None + if entity_key_hex: + try: + entity_key_bytes = bytes.fromhex(entity_key_hex) + entity_key_proto = deserialize_entity_key(entity_key_bytes) + except Exception: + pass + + res: Dict[str, ValueProto] = {} + for feature_name in requested_features: + if feature_name in vector_fields and values: + val = _serialize_vector_to_float_list(embedding) + res[feature_name] = val + elif feature_name in entity_name_type_map: + from feast.types import PrimitiveFeastType + + feat_type = entity_name_type_map[feature_name] + if feat_type == PrimitiveFeastType.STRING: + raw_val = metadata.get(feature_name, "") + res[feature_name] = ValueProto(string_val=str(raw_val)) + elif feat_type in ( + PrimitiveFeastType.INT64, + PrimitiveFeastType.INT32, + ): + raw_val = metadata.get(feature_name, 0) + res[feature_name] = ValueProto(int64_val=int(raw_val)) + elif feat_type == PrimitiveFeastType.BYTES: + raw_val = metadata.get(feature_name, "") + try: + res[feature_name] = ValueProto( + bytes_val=base64.b64decode(raw_val) + ) + except Exception: + res[feature_name] = ValueProto(string_val=str(raw_val)) + elif feature_name in metadata: + val = _metadata_to_proto_value(metadata[feature_name], None) + res[feature_name] = val + + res["distance"] = ValueProto(float_val=score) + result_list.append((res_ts, entity_key_proto, res if res else None)) + + return result_list + + def _delete_namespace(self, config: RepoConfig, table: FeatureView) -> None: + """Delete all vectors in the namespace for the given feature view.""" + try: + index = self._get_index(config) + namespace = self._get_namespace(config, table) + index.delete(delete_all=True, namespace=namespace) + except Exception: + logger.exception("Error deleting namespace for feature view %s", table.name) + + @staticmethod + def _wait_for_index_ready(client: Any, index_name: str) -> None: + deadline = time.time() + _INDEX_READY_TIMEOUT + while time.time() < deadline: + desc = client.describe_index(index_name) + status = getattr(desc, "status", {}) + if isinstance(status, dict): + ready = status.get("ready", False) + else: + ready = getattr(status, "ready", False) + if ready: + return + time.sleep(2) + logger.warning( + "Pinecone index '%s' did not become ready within %ds", + index_name, + _INDEX_READY_TIMEOUT, + ) + + +def _table_id(project: str, table: FeatureView, enable_versioning: bool = False) -> str: + return compute_table_id(project, table, enable_versioning) + + +def _extract_vector(value_proto: ValueProto) -> Optional[List[float]]: + """Extract a float list from a ValueProto that contains a vector.""" + if value_proto.HasField("float_list_val"): + return list(value_proto.float_list_val.val) + if value_proto.HasField("double_list_val"): + return [float(v) for v in value_proto.double_list_val.val] + if value_proto.HasField("int32_list_val"): + return [float(v) for v in value_proto.int32_list_val.val] + if value_proto.HasField("int64_list_val"): + return [float(v) for v in value_proto.int64_list_val.val] + return None + + +def _proto_value_to_metadata(value_proto: ValueProto) -> Any: + """Convert a ValueProto to a JSON-serializable value for Pinecone metadata.""" + if value_proto.HasField("string_val"): + return value_proto.string_val + if value_proto.HasField("int32_val"): + return value_proto.int32_val + if value_proto.HasField("int64_val"): + return value_proto.int64_val + if value_proto.HasField("float_val"): + return value_proto.float_val + if value_proto.HasField("double_val"): + return value_proto.double_val + if value_proto.HasField("bool_val"): + return value_proto.bool_val + + if value_proto.HasField("float_list_val"): + return list(value_proto.float_list_val.val) + if value_proto.HasField("double_list_val"): + return list(value_proto.double_list_val.val) + if value_proto.HasField("int32_list_val"): + return list(value_proto.int32_list_val.val) + if value_proto.HasField("int64_list_val"): + return list(value_proto.int64_list_val.val) + + if value_proto.HasField("bytes_val"): + return base64.b64encode(value_proto.bytes_val).decode("utf-8") + + return base64.b64encode(value_proto.SerializeToString()).decode("utf-8") + + +def _metadata_to_proto_value(metadata_value: Any, feast_type: Any) -> ValueProto: + """Convert a Pinecone metadata value back to a ValueProto.""" + val = ValueProto() + + if feast_type is not None: + from feast.type_map import VALUE_TYPE_TO_PROTO_VALUE_MAP + from feast.types import from_feast_type + + value_type = from_feast_type(feast_type) + proto_attr = VALUE_TYPE_TO_PROTO_VALUE_MAP.get(value_type) + + if proto_attr: + if proto_attr in ("int32_val", "int64_val"): + setattr(val, proto_attr, int(metadata_value)) + return val + elif proto_attr in ("float_val", "double_val"): + setattr(val, proto_attr, float(metadata_value)) + return val + elif proto_attr == "bool_val": + setattr(val, proto_attr, bool(metadata_value)) + return val + elif proto_attr == "string_val": + setattr(val, proto_attr, str(metadata_value)) + return val + elif proto_attr == "bytes_val": + if isinstance(metadata_value, str): + setattr(val, proto_attr, base64.b64decode(metadata_value)) + else: + setattr(val, proto_attr, metadata_value) + return val + elif proto_attr in ( + "float_list_val", + "double_list_val", + "int32_list_val", + "int64_list_val", + ): + if isinstance(metadata_value, list): + getattr(val, proto_attr).val.extend(metadata_value) + return val + + if isinstance(metadata_value, bool): + val.bool_val = metadata_value + elif isinstance(metadata_value, int): + val.int64_val = metadata_value + elif isinstance(metadata_value, float): + val.double_val = metadata_value + elif isinstance(metadata_value, str): + val.string_val = metadata_value + elif isinstance(metadata_value, list): + if metadata_value and isinstance(metadata_value[0], (int, float)): + val.float_list_val.val.extend([float(v) for v in metadata_value]) + else: + val.string_val = str(metadata_value) + else: + val.string_val = str(metadata_value) + + return val diff --git a/sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone_repo_configuration.py b/sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone_repo_configuration.py new file mode 100644 index 00000000000..483c2a8c1b3 --- /dev/null +++ b/sdk/python/feast/infra/online_stores/pinecone_online_store/pinecone_repo_configuration.py @@ -0,0 +1,12 @@ +from tests.universal.feature_repos.integration_test_repo_config import ( + IntegrationTestRepoConfig, +) +from tests.universal.feature_repos.universal.online_store.pinecone import ( + PineconeOnlineStoreCreator, +) + +FULL_REPO_CONFIGS = [ + IntegrationTestRepoConfig( + online_store="pinecone", online_store_creator=PineconeOnlineStoreCreator + ), +] diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index de83a2db163..37a4477e3bf 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -86,6 +86,7 @@ "qdrant": "feast.infra.online_stores.qdrant_online_store.qdrant.QdrantOnlineStore", "couchbase.online": "feast.infra.online_stores.couchbase_online_store.couchbase.CouchbaseOnlineStore", "milvus": "feast.infra.online_stores.milvus_online_store.milvus.MilvusOnlineStore", + "pinecone": "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore", "mongodb": "feast.infra.online_stores.mongodb_online_store.MongoDBOnlineStore", "hybrid": "feast.infra.online_stores.hybrid_online_store.hybrid_online_store.HybridOnlineStore", **LEGACY_ONLINE_STORE_CLASS_FOR_TYPE, diff --git a/sdk/python/tests/unit/infra/online_store/test_pinecone.py b/sdk/python/tests/unit/infra/online_store/test_pinecone.py new file mode 100644 index 00000000000..fc3fb87c1a2 --- /dev/null +++ b/sdk/python/tests/unit/infra/online_store/test_pinecone.py @@ -0,0 +1,518 @@ +"""Unit tests for the Pinecone online store.""" + +from datetime import datetime, timedelta +from unittest.mock import MagicMock, patch + +import pytest + +from feast import Entity, FeatureView +from feast.field import Field +from feast.infra.key_encoding_utils import serialize_entity_key +from feast.infra.online_stores.pinecone_online_store.pinecone import ( + PineconeOnlineStore, + PineconeOnlineStoreConfig, + _extract_vector, + _metadata_to_proto_value, + _proto_value_to_metadata, + _table_id, +) +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.types import Array, Float32, Int64, String +from feast.value_type import ValueType + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_feature_view(name="driver_stats", version_number=None, version_tag=None): + entity = Entity( + name="driver_id", + join_keys=["driver_id"], + value_type=ValueType.INT64, + ) + fv = FeatureView( + name=name, + entities=[entity], + ttl=timedelta(days=1), + schema=[ + Field(name="trips_today", dtype=Float32), + Field(name="driver_name", dtype=String), + ], + ) + if version_number is not None: + fv.current_version_number = version_number + if version_tag is not None: + fv.projection.version_tag = version_tag + return fv + + +def _make_vector_feature_view(): + entity = Entity( + name="item_id", + join_keys=["item_id"], + value_type=ValueType.INT64, + ) + return FeatureView( + name="embedded_docs", + entities=[entity], + ttl=timedelta(days=1), + schema=[ + Field( + name="vector", + dtype=Array(Float32), + vector_index=True, + vector_search_metric="COSINE", + ), + Field(name="item_id", dtype=Int64), + Field(name="sentence_chunks", dtype=String), + ], + ) + + +def _make_config(project="test_project", versioning=False): + config = MagicMock() + config.project = project + config.entity_key_serialization_version = 3 + config.registry.enable_online_feature_view_versioning = versioning + config.online_store = PineconeOnlineStoreConfig( + type="pinecone", + api_key="test-api-key", # pragma: allowlist secret + index_name="test-index", + embedding_dim=3, + metric="cosine", + vector_enabled=True, + ) + return config + + +def _make_entity_key(driver_id: int = 1) -> EntityKeyProto: + key = EntityKeyProto() + key.join_keys.append("driver_id") + val = ValueProto(int64_val=driver_id) + key.entity_values.append(val) + return key + + +# --------------------------------------------------------------------------- +# Config Tests +# --------------------------------------------------------------------------- + + +class TestPineconeOnlineStoreConfig: + def test_defaults(self): + config = PineconeOnlineStoreConfig() + assert config.type == "pinecone" + assert config.index_name == "feast-online" + assert config.embedding_dim == 128 + assert config.metric == "cosine" + assert config.cloud == "aws" + assert config.region == "us-east-1" + assert config.vector_enabled is True + + def test_custom_values(self): + config = PineconeOnlineStoreConfig( + api_key="my-key", # pragma: allowlist secret + index_name="my-index", + embedding_dim=384, + metric="dotproduct", + cloud="gcp", + region="us-central1", + namespace="custom-ns", + ) + assert config.api_key == "my-key" # pragma: allowlist secret + assert config.index_name == "my-index" + assert config.embedding_dim == 384 + assert config.metric == "dotproduct" + assert config.cloud == "gcp" + assert config.region == "us-central1" + assert config.namespace == "custom-ns" + + +# --------------------------------------------------------------------------- +# Table ID Tests +# --------------------------------------------------------------------------- + + +class TestTableId: + def test_no_versioning(self): + fv = _make_feature_view() + assert _table_id("test_project", fv) == "test_project_driver_stats" + + def test_versioning_enabled_with_version(self): + fv = _make_feature_view(version_number=2) + assert ( + _table_id("test_project", fv, enable_versioning=True) + == "test_project_driver_stats_v2" + ) + + def test_versioning_disabled_ignores_version(self): + fv = _make_feature_view(version_number=5) + assert _table_id("test_project", fv) == "test_project_driver_stats" + + +# --------------------------------------------------------------------------- +# Proto Conversion Tests +# --------------------------------------------------------------------------- + + +class TestProtoConversions: + def test_extract_vector_float_list(self): + val = ValueProto() + val.float_list_val.val.extend([1.0, 2.0, 3.0]) + result = _extract_vector(val) + assert result == [1.0, 2.0, 3.0] + + def test_extract_vector_double_list(self): + val = ValueProto() + val.double_list_val.val.extend([1.0, 2.0]) + result = _extract_vector(val) + assert result == [1.0, 2.0] + + def test_extract_vector_none_for_non_list(self): + val = ValueProto(string_val="hello") + assert _extract_vector(val) is None + + def test_proto_value_to_metadata_string(self): + val = ValueProto(string_val="hello") + assert _proto_value_to_metadata(val) == "hello" + + def test_proto_value_to_metadata_int(self): + val = ValueProto(int64_val=42) + assert _proto_value_to_metadata(val) == 42 + + def test_proto_value_to_metadata_float(self): + val = ValueProto(float_val=3.14) + assert abs(_proto_value_to_metadata(val) - 3.14) < 1e-6 + + def test_proto_value_to_metadata_bool(self): + val = ValueProto(bool_val=True) + assert _proto_value_to_metadata(val) is True + + def test_metadata_to_proto_value_string(self): + result = _metadata_to_proto_value("hello", None) + assert result.string_val == "hello" + + def test_metadata_to_proto_value_int(self): + result = _metadata_to_proto_value(42, None) + assert result.int64_val == 42 + + def test_metadata_to_proto_value_float(self): + result = _metadata_to_proto_value(3.14, None) + assert abs(result.double_val - 3.14) < 1e-6 + + def test_metadata_to_proto_value_bool(self): + result = _metadata_to_proto_value(True, None) + assert result.bool_val is True + + def test_metadata_to_proto_value_list(self): + result = _metadata_to_proto_value([1.0, 2.0, 3.0], None) + assert list(result.float_list_val.val) == [1.0, 2.0, 3.0] + + +# --------------------------------------------------------------------------- +# Online Store Tests (mocked Pinecone client) +# --------------------------------------------------------------------------- + + +class TestPineconeOnlineStoreWriteBatch: + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_online_write_batch(self, mock_get_index): + mock_index = MagicMock() + mock_get_index.return_value = mock_index + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_feature_view() + entity_key = _make_entity_key(1) + + data = [ + ( + entity_key, + { + "trips_today": ValueProto(float_val=10.0), + "driver_name": ValueProto(string_val="Alice"), + }, + datetime(2024, 1, 1, 12, 0, 0), + None, + ) + ] + + store.online_write_batch(config, fv, data, progress=None) + + mock_index.upsert.assert_called_once() + call_kwargs = mock_index.upsert.call_args + vectors = call_kwargs.kwargs.get("vectors") or call_kwargs[1].get("vectors") + assert len(vectors) == 1 + vec = vectors[0] + assert "id" in vec + assert "values" in vec + assert "metadata" in vec + assert vec["metadata"]["trips_today"] == 10.0 + assert vec["metadata"]["driver_name"] == "Alice" + + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_write_batch_deduplication(self, mock_get_index): + mock_index = MagicMock() + mock_get_index.return_value = mock_index + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_feature_view() + entity_key = _make_entity_key(1) + + data = [ + ( + entity_key, + {"trips_today": ValueProto(float_val=5.0)}, + datetime(2024, 1, 1, 10, 0, 0), + None, + ), + ( + entity_key, + {"trips_today": ValueProto(float_val=15.0)}, + datetime(2024, 1, 1, 14, 0, 0), + None, + ), + ] + + store.online_write_batch(config, fv, data, progress=None) + + call_kwargs = mock_index.upsert.call_args + vectors = call_kwargs.kwargs.get("vectors") or call_kwargs[1].get("vectors") + assert len(vectors) == 1 + assert vectors[0]["metadata"]["trips_today"] == 15.0 + + +class TestPineconeOnlineStoreRead: + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_online_read(self, mock_get_index): + mock_index = MagicMock() + mock_get_index.return_value = mock_index + + entity_key = _make_entity_key(1) + entity_key_str = serialize_entity_key( + entity_key, entity_key_serialization_version=3 + ).hex() + + now_ts = int(datetime(2024, 1, 1, 12, 0, 0).timestamp() * 1e6) + + mock_record = MagicMock() + mock_record.metadata = { + "event_ts": now_ts, + "created_ts": 0, + "entity_key": entity_key_str, + "trips_today": 10.0, + "driver_name": "Alice", + } + mock_record.values = [0.0, 0.0, 0.0] + + mock_response = MagicMock() + mock_response.vectors = {entity_key_str: mock_record} + mock_index.fetch.return_value = mock_response + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_feature_view() + + result = store.online_read( + config, fv, [entity_key], requested_features=["trips_today"] + ) + + assert len(result) == 1 + ts, features = result[0] + assert ts is not None + assert features is not None + assert "trips_today" in features + + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_online_read_missing_entity(self, mock_get_index): + mock_index = MagicMock() + mock_get_index.return_value = mock_index + + mock_response = MagicMock() + mock_response.vectors = {} + mock_index.fetch.return_value = mock_response + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_feature_view() + entity_key = _make_entity_key(999) + + result = store.online_read(config, fv, [entity_key]) + + assert len(result) == 1 + assert result[0] == (None, None) + + +class TestPineconeOnlineStoreUpdate: + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_client" + ) + def test_update_creates_index(self, mock_get_client): + import sys + import types + + mock_pinecone_module = types.ModuleType("pinecone") + mock_pinecone_module.ServerlessSpec = MagicMock() + sys.modules["pinecone"] = mock_pinecone_module + + try: + mock_client = MagicMock() + mock_get_client.return_value = mock_client + mock_client.list_indexes.return_value = [] + + mock_desc = MagicMock() + mock_desc.status = {"ready": True} + mock_client.describe_index.return_value = mock_desc + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_feature_view() + + store.update(config, [], [fv], [], [], partial=False) + + mock_client.create_index.assert_called_once() + call_kwargs = mock_client.create_index.call_args.kwargs + assert call_kwargs["name"] == "test-index" + assert call_kwargs["dimension"] == 3 + assert call_kwargs["metric"] == "cosine" + finally: + del sys.modules["pinecone"] + + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_client" + ) + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_update_skips_existing_index(self, mock_get_index, mock_get_client): + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + existing_idx = MagicMock() + existing_idx.name = "test-index" + mock_client.list_indexes.return_value = [existing_idx] + mock_get_index.return_value = MagicMock() + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_feature_view() + + store.update(config, [], [fv], [], [], partial=False) + + mock_client.create_index.assert_not_called() + + +class TestPineconeOnlineStoreTeardown: + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_teardown_deletes_namespace(self, mock_get_index): + mock_index = MagicMock() + mock_get_index.return_value = mock_index + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_feature_view() + + store.teardown(config, [fv], []) + + mock_index.delete.assert_called_once() + call_kwargs = mock_index.delete.call_args.kwargs + assert call_kwargs["delete_all"] is True + assert call_kwargs["namespace"] == "test_project_driver_stats" + + +class TestPineconeRetrieveDocumentsV2: + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_retrieve_online_documents_v2(self, mock_get_index): + mock_index = MagicMock() + mock_get_index.return_value = mock_index + + entity_key = EntityKeyProto() + entity_key.join_keys.append("item_id") + entity_key.entity_values.append(ValueProto(int64_val=42)) + entity_key_str = serialize_entity_key( + entity_key, entity_key_serialization_version=3 + ).hex() + + now_ts = int(datetime(2024, 1, 1, 12, 0, 0).timestamp() * 1e6) + + mock_match = MagicMock() + mock_match.metadata = { + "event_ts": now_ts, + "entity_key": entity_key_str, + "sentence_chunks": "New York City", + } + mock_match.values = [0.1, 0.2, 0.3] + mock_match.score = 0.95 + + mock_response = MagicMock() + mock_response.matches = [mock_match] + mock_index.query.return_value = mock_response + + store = PineconeOnlineStore() + config = _make_config() + fv = _make_vector_feature_view() + + results = store.retrieve_online_documents_v2( + config=config, + table=fv, + requested_features=["vector", "sentence_chunks"], + embedding=[0.1, 0.2, 0.3], + top_k=3, + distance_metric="cosine", + ) + + assert len(results) == 1 + ts, entity, features = results[0] + assert ts is not None + assert features is not None + assert "distance" in features + assert features["distance"].float_val == pytest.approx(0.95) + assert "sentence_chunks" in features + + @patch( + "feast.infra.online_stores.pinecone_online_store.pinecone.PineconeOnlineStore._get_index" + ) + def test_retrieve_requires_embedding(self, mock_get_index): + store = PineconeOnlineStore() + config = _make_config() + fv = _make_vector_feature_view() + + with pytest.raises(ValueError, match="requires a query embedding"): + store.retrieve_online_documents_v2( + config=config, + table=fv, + requested_features=["vector"], + embedding=None, + top_k=3, + query_string="some text", + ) + + def test_retrieve_vector_not_enabled(self): + store = PineconeOnlineStore() + config = _make_config() + config.online_store.vector_enabled = False + fv = _make_vector_feature_view() + + with pytest.raises(ValueError, match="not enabled"): + store.retrieve_online_documents_v2( + config=config, + table=fv, + requested_features=["vector"], + embedding=[0.1, 0.2, 0.3], + top_k=3, + ) diff --git a/sdk/python/tests/universal/feature_repos/universal/online_store/pinecone.py b/sdk/python/tests/universal/feature_repos/universal/online_store/pinecone.py new file mode 100644 index 00000000000..04dd9dd481f --- /dev/null +++ b/sdk/python/tests/universal/feature_repos/universal/online_store/pinecone.py @@ -0,0 +1,27 @@ +import os +from typing import Any + +from tests.universal.feature_repos.universal.online_store_creator import ( + OnlineStoreCreator, +) + + +class PineconeOnlineStoreCreator(OnlineStoreCreator): + def __init__(self, project_name: str, **kwargs): + super().__init__(project_name) + + def create_online_store(self) -> dict[str, Any]: + api_key = os.environ.get("PINECONE_API_KEY", "") + return { + "type": "pinecone", + "api_key": api_key, + "index_name": f"feast-test-{self.project_name}", + "embedding_dim": 2, + "metric": "cosine", + "vector_enabled": True, + "cloud": "aws", + "region": "us-east-1", + } + + def teardown(self): + pass