Lee-W commented on code in PR #37307:
URL: https://github.com/apache/airflow/pull/37307#discussion_r1570128288
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -135,119 +170,132 @@ def upsert(
**kwargs,
)
- @staticmethod
+ def get_pod_spec_obj(
+ self,
+ replicas: int | None = None,
+ shards: int | None = None,
+ pods: int | None = None,
+ pod_type: str | None = "p1.x1",
+ metadata_config: dict | None = None,
+ source_collection: str | None = None,
+ environment: str | None = None,
+ ) -> PodSpec:
+ """
+ Get a PodSpec object.
+
+ :param replicas: The number of replicas.
+ :param shards: The number of shards.
+ :param pods: The number of pods.
+ :param pod_type: The type of pod.
+ :param metadata_config: The metadata configuration.
+ :param source_collection: The source collection.
+ :param environment: The environment to use when creating the index.
+ """
+ return PodSpec(
+ environment=environment or self.environment,
+ replicas=replicas,
+ shards=shards,
+ pods=pods,
+ pod_type=pod_type,
+ metadata_config=metadata_config,
+ source_collection=source_collection,
+ )
+
+ def get_serverless_spec_obj(self, cloud, region: str | None = None) ->
ServerlessSpec:
Review Comment:
```suggestion
def get_serverless_spec_obj(self, cloud: str, region: str | None = None)
-> ServerlessSpec:
```
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -135,119 +170,132 @@ def upsert(
**kwargs,
)
- @staticmethod
+ def get_pod_spec_obj(
+ self,
+ replicas: int | None = None,
Review Comment:
nitpick: does the order of these arguments matter? if not, should we make
them keyword arguments
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -135,119 +170,132 @@ def upsert(
**kwargs,
)
- @staticmethod
+ def get_pod_spec_obj(
+ self,
+ replicas: int | None = None,
+ shards: int | None = None,
+ pods: int | None = None,
+ pod_type: str | None = "p1.x1",
+ metadata_config: dict | None = None,
+ source_collection: str | None = None,
+ environment: str | None = None,
+ ) -> PodSpec:
+ """
+ Get a PodSpec object.
+
+ :param replicas: The number of replicas.
+ :param shards: The number of shards.
+ :param pods: The number of pods.
+ :param pod_type: The type of pod.
+ :param metadata_config: The metadata configuration.
+ :param source_collection: The source collection.
+ :param environment: The environment to use when creating the index.
+ """
+ return PodSpec(
+ environment=environment or self.environment,
+ replicas=replicas,
+ shards=shards,
+ pods=pods,
+ pod_type=pod_type,
+ metadata_config=metadata_config,
+ source_collection=source_collection,
+ )
+
+ def get_serverless_spec_obj(self, cloud, region: str | None = None) ->
ServerlessSpec:
+ """
+ Get a ServerlessSpec object.
+
+ :param cloud: The cloud provider.
+ :param region: The region to use when creating the index.
+ """
+ return ServerlessSpec(cloud=cloud, region=region or self.region)
+
def create_index(
+ self,
index_name: str,
dimension: int,
- index_type: str | None = "approximated",
+ spec: ServerlessSpec | PodSpec,
metric: str | None = "cosine",
- replicas: int | None = 1,
- shards: int | None = 1,
- pods: int | None = 1,
- pod_type: str | None = "p1",
- index_config: dict[str, str] | None = None,
- metadata_config: dict[str, str] | None = None,
- source_collection: str | None = "",
timeout: int | None = None,
) -> None:
"""
Create a new index.
- .. seealso:: https://docs.pinecone.io/reference/create_index/
-
- :param index_name: The name of the index to create.
- :param dimension: the dimension of vectors that would be inserted in
the index
- :param index_type: type of index, one of {"approximated", "exact"},
defaults to "approximated".
- :param metric: type of metric used in the vector index, one of
{"cosine", "dotproduct", "euclidean"}
- :param replicas: the number of replicas, defaults to 1.
- :param shards: the number of shards per index, defaults to 1.
- :param pods: Total number of pods to be used by the index. pods =
shard*replicas
- :param pod_type: the pod type to be used for the index. can be one of
p1 or s1.
- :param index_config: Advanced configuration options for the index
- :param metadata_config: Configuration related to the metadata index
- :param source_collection: Collection name to create the index from
- :param timeout: Timeout for wait until index gets ready.
+ :param index_name: The name of the index.
+ :param dimension: The dimension of the vectors to be indexed.
+ :param spec: Pass a `ServerlessSpec` object to create a serverless
index or a `PodSpec` object to create a pod index.
+ :param metric: The metric to use.
+ :param timeout: The timeout to use.
"""
- pinecone.create_index(
+ self.pc.create_index(
name=index_name,
- timeout=timeout,
- index_type=index_type,
dimension=dimension,
+ spec=spec,
metric=metric,
- pods=pods,
- replicas=replicas,
- shards=shards,
- pod_type=pod_type,
- metadata_config=metadata_config,
- source_collection=source_collection,
- index_config=index_config,
+ timeout=timeout,
)
- @staticmethod
- def describe_index(index_name: str) -> Any:
+ def describe_index(self, index_name: str) -> Any:
"""
Retrieve information about a specific index.
:param index_name: The name of the index to describe.
"""
- return pinecone.describe_index(name=index_name)
+ return self.pc.describe_index(name=index_name)
- @staticmethod
- def delete_index(index_name: str, timeout: int | None = None) -> None:
+ def delete_index(self, index_name: str, timeout: int | None = None) ->
None:
"""
Delete a specific index.
:param index_name: the name of the index.
:param timeout: Timeout for wait until index gets ready.
"""
- pinecone.delete_index(name=index_name, timeout=timeout)
+ self.pc.delete_index(name=index_name, timeout=timeout)
- @staticmethod
- def configure_index(index_name: str, replicas: int | None = None,
pod_type: str | None = "") -> None:
+ def configure_index(
+ self, index_name: str, replicas: int | None = None, pod_type: str |
None = ""
+ ) -> None:
"""
Change the current configuration of the index.
:param index_name: The name of the index to configure.
:param replicas: The new number of replicas.
:param pod_type: the new pod_type for the index.
"""
- pinecone.configure_index(name=index_name, replicas=replicas,
pod_type=pod_type)
+ self.pc.configure_index(name=index_name, replicas=replicas,
pod_type=pod_type)
@staticmethod
Review Comment:
```suggestion
```
I guess we probably should remove it as well?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
