You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "sunank200 (via GitHub)" <gi...@apache.org> on 2023/10/12 13:33:41 UTC
[PR] Enable pgvector support for Postgres provider [airflow]
sunank200 opened a new pull request, #34891:
URL: https://github.com/apache/airflow/pull/34891
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
This PR is part of our larger effort to add first-class integrations to support LLMOps that was presented at the Airflow Summit. This PR specifically adds the [pgvector](https://github.com/pgvector/pgvector) support for Postgres Provider. [pgvector](https://github.com/pgvector/pgvector#pgvector) is a renowned Open-source vector similarity search for Postgres. In this iteration, we are integrating with their Embeddings Model.
The primary objective of this Provider is to present users with an alternative embedding model. This allows them to generate vectors for their proprietary data, a pivotal step towards establishing integrations with LLM models like ChatGPT.
Example DAG:
The `PgVectorIngestOperator` can accept either a list of strings or a callable returning a list of strings.
```
```
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1358562950
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -82,3 +87,61 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
+ """
+ Operator for ingesting text and embeddings into a PostgreSQL database using the pgvector library.
+
+ :param conn_id: The connection ID for the postgresql database.
+ :param input_text: The input text to be ingested.
+ :param input_embedding: The input embedding associated with the text, provided as a list of floats.
+ :param input_callable: A callable that returns the embedding if 'input_embedding' is not provided.
+ :param input_callable_args: Positional arguments for the 'input_callable'.
+ :param input_callable_kwargs: Keyword arguments for the 'input_callable'.
+ :param kwargs: Additional keyword arguments for the BaseOperator.
+ """
+
+ def __init__(
+ self,
+ conn_id: str,
+ input_text: str | None = None,
+ input_embedding: list[float] | None = None,
+ input_callable: Callable[[Any], Any] | None = None,
+ input_callable_args: Collection[Any] | None = None,
+ input_callable_kwargs: Mapping[str, Any] | None = None,
+ **kwargs: Any,
+ ) -> None:
+ self.conn_id = conn_id
+ self.input_text = input_text
+ self.input_embedding = input_embedding
+ self.input_callable = input_callable
+ self.input_callable_args = input_callable_args
+ self.input_callable_kwargs = input_callable_kwargs
+ super().__init__(**kwargs)
+
+ def execute(self, context: Context) -> Any:
+ """
+ Executes the ingestion process.
+
+ This method either uses the provided embedding or computes the embedding using the
+ provided callable. The text and its associated embedding are then stored into the
+ PostgreSQL database's 'documents' table.
+
+ :param context: The execution context for the operator.
+ """
+ from pgvector.psycopg import register_vector
+
+ if all([self.input_embedding, self.input_callable]):
+ raise AirflowException("Only one of 'input_embedding' and 'input_callable' is allowed")
Review Comment:
Looks like expected here `airflow.utils.helpers.exactly_one`:
https://github.com/apache/airflow/blob/58d8577f3528e7e143addefb03711bdc82cb1782/airflow/utils/helpers.py#L297-L307
Some example of usage
https://github.com/apache/airflow/blob/0c8e30e43b70e9d033e1686b327eb00aab82479c/airflow/providers/slack/hooks/slack.py#L190-L192
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "sunank200 (via GitHub)" <gi...@apache.org>.
sunank200 commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1357077212
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,29 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(
+ self, table: str, input_data: list[str], embeddings: list[float], vector_size: int
+ ) -> None:
+ """
+ Store embedding vector in Postgres table.
+
+ :param table: The Name of the table
+ :param input_data: The source data from which the embedding has been created
+ :param embeddings: The embedding vector response from LLM service
+ :param vector_size: The size of vector. The maximum dimensions can be 2,000
+ """
+ from pgvector.psycopg import register_vector
+
+ self.conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+ register_vector(self.conn)
+ self.conn.execute(
+ "CREATE TABLE IF NOT EXISTS %s (id bigserial PRIMARY KEY, content text, embedding vector(%s))",
+ (table, vector_size),
+ )
+
+ for content, embedding in zip(input_data, embeddings):
+ self.execute("INSERT INTO %s (content, embedding) VALUES (%s, %s)", (table, content, embedding))
Review Comment:
https://github.com/apache/airflow/pull/34891/commits/90dfdba912eaf82795c234a0bae6db2f1e0ac175
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -82,3 +85,42 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
Review Comment:
https://github.com/apache/airflow/pull/34891/commits/90dfdba912eaf82795c234a0bae6db2f1e0ac175
##########
airflow/providers/postgres/provider.yaml:
##########
@@ -53,6 +53,7 @@ dependencies:
- apache-airflow>=2.4.0
- apache-airflow-providers-common-sql>=1.3.1
- psycopg2-binary>=2.8.0
+ - pgvector
Review Comment:
https://github.com/apache/airflow/pull/34891/commits/90dfdba912eaf82795c234a0bae6db2f1e0ac175
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "jlaneve (via GitHub)" <gi...@apache.org>.
jlaneve commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1363762415
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -80,3 +86,60 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
+ """
+ Operator for ingesting text and embeddings into a PostgreSQL database using the pgvector library.
+
+ :param conn_id: The connection ID for the postgresql database.
+ :param input_data: Tuple containing the string input content and corresponding list of float vector
+ embeddings.
+ :param input_callable: A callable that returns a tuple containing the string input content and
+ corresponding list of float vector embeddings, if ``input_data`` is not provided.
+ :param input_callable_args: Positional arguments for the 'input_callable'.
+ :param input_callable_kwargs: Keyword arguments for the 'input_callable'.
+ :param kwargs: Additional keyword arguments for the BaseOperator.
+ """
+
+ def __init__(
+ self,
+ conn_id: str,
+ input_data: tuple[str, list[float]] | None = None,
+ input_callable: Callable[[Any], Any] | None = None,
+ input_callable_args: Collection[Any] | None = None,
+ input_callable_kwargs: Mapping[str, Any] | None = None,
Review Comment:
yes, purpose here is to give a workaround to storing things in XComs. depending on how much data you're working with, passing a large number of vectors through XComs can be unideal (especially if you don't have a custom XCom backend). instead, giving the user the ability to execute the same data fetching code within the task means we don't pollute XComs.
this has disadvantages though, particularly around retries. hence why there are two different input methods to let the user decide which is right for them: (1) from XComs and (2) with a callable
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1356845699
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -82,3 +85,42 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
Review Comment:
Needs docs
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on PR #34891:
URL: https://github.com/apache/airflow/pull/34891#issuecomment-1792174444
Closing this PR in lieu of a separate provider PR: https://github.com/apache/airflow/pull/35399
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti closed pull request #34891: Enable pgvector support for Postgres provider
URL: https://github.com/apache/airflow/pull/34891
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1364142683
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -80,3 +86,60 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
+ """
+ Operator for ingesting text and embeddings into a PostgreSQL database using the pgvector library.
+
+ :param conn_id: The connection ID for the postgresql database.
+ :param input_data: Tuple containing the string input content and corresponding list of float vector
+ embeddings.
+ :param input_callable: A callable that returns a tuple containing the string input content and
+ corresponding list of float vector embeddings, if ``input_data`` is not provided.
+ :param input_callable_args: Positional arguments for the 'input_callable'.
+ :param input_callable_kwargs: Keyword arguments for the 'input_callable'.
+ :param kwargs: Additional keyword arguments for the BaseOperator.
+ """
+
+ def __init__(
+ self,
+ conn_id: str,
+ input_data: tuple[str, list[float]] | None = None,
+ input_callable: Callable[[Any], Any] | None = None,
+ input_callable_args: Collection[Any] | None = None,
+ input_callable_kwargs: Mapping[str, Any] | None = None,
Review Comment:
> (2) with a callable
That is exactly what a PythonOperator (and taskflow decorators) does as well as provide greater flexibility, e.g. access to task context.
In the other hand current implementation it is a combination of (1) classic operator and (2) restricted PythonOperator. The second part could be replaced by add examples in docs how to use Hook within the taskflow.
It would reduce complexibility of the code and number of required tests
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1363167734
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -80,3 +86,60 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
+ """
+ Operator for ingesting text and embeddings into a PostgreSQL database using the pgvector library.
+
+ :param conn_id: The connection ID for the postgresql database.
+ :param input_data: Tuple containing the string input content and corresponding list of float vector
+ embeddings.
+ :param input_callable: A callable that returns a tuple containing the string input content and
+ corresponding list of float vector embeddings, if ``input_data`` is not provided.
+ :param input_callable_args: Positional arguments for the 'input_callable'.
+ :param input_callable_kwargs: Keyword arguments for the 'input_callable'.
+ :param kwargs: Additional keyword arguments for the BaseOperator.
+ """
+
+ def __init__(
+ self,
+ conn_id: str,
+ input_data: tuple[str, list[float]] | None = None,
+ input_callable: Callable[[Any], Any] | None = None,
+ input_callable_args: Collection[Any] | None = None,
+ input_callable_kwargs: Mapping[str, Any] | None = None,
Review Comment:
One of the approaches we would like to support within the operator is to get data without using XCOMs, because if I am right there is no automatic cleanup of XCOMs, so users might not want to populate the database with large number of XCOM entries. Hence, we thought of supporting the callable approach within the operator
cc: @jlaneve
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1362805154
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,34 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(self, table: str, input_data: Iterable[tuple[str, float]], vector_size: int) -> None:
Review Comment:
```suggestion
def ingest_embedding(self, table: str, input_data: Iterable[tuple[str, list[float]]], vector_size: int) -> None:
```
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1363227895
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,34 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(self, table: str, input_data: Iterable[tuple[str, list[float]]], vector_size: int) -> None:
+ """
+ Store embedding vector in Postgres table.
+
+ :param table: The Name of the table
+ :param input_data: Iterable containing tuples of input data and corresponding embedding vectors.
+ :param vector_size: The size of vector. The maximum dimensions can be 2,000
+ """
+ from pgvector.psycopg import register_vector
+ from psycopg2 import sql
+
+ self.conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+ register_vector(self.conn)
+
+ create_table_query = sql.SQL(
+ "CREATE TABLE IF NOT EXISTS {} (id bigserial PRIMARY KEY, content text, embedding vector({}))"
+ ).format(sql.Identifier(table), sql.Literal(vector_size))
+
+ self.conn.execute(create_table_query)
+
+ for data_item in input_data:
+ insert_query = sql.SQL("INSERT INTO {} (content, embedding) VALUES (%s, %s)").format(
+ sql.Identifier(table)
+ )
+ content = data_item[0]
+ embedding = data_item[1]
+ self.conn.execute(insert_query, (content, embedding))
Review Comment:
```suggestion
self.conn.execute(insert_query, data)
```
The unpacking seems unnecessary since the two are joined immediately?
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,34 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(self, table: str, input_data: Iterable[tuple[str, list[float]]], vector_size: int) -> None:
+ """
+ Store embedding vector in Postgres table.
+
+ :param table: The Name of the table
+ :param input_data: Iterable containing tuples of input data and corresponding embedding vectors.
+ :param vector_size: The size of vector. The maximum dimensions can be 2,000
+ """
+ from pgvector.psycopg import register_vector
+ from psycopg2 import sql
+
+ self.conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+ register_vector(self.conn)
+
+ create_table_query = sql.SQL(
+ "CREATE TABLE IF NOT EXISTS {} (id bigserial PRIMARY KEY, content text, embedding vector({}))"
+ ).format(sql.Identifier(table), sql.Literal(vector_size))
+
+ self.conn.execute(create_table_query)
+
+ for data_item in input_data:
+ insert_query = sql.SQL("INSERT INTO {} (content, embedding) VALUES (%s, %s)").format(
+ sql.Identifier(table)
+ )
+ content = data_item[0]
+ embedding = data_item[1]
+ self.conn.execute(insert_query, (content, embedding))
Review Comment:
```suggestion
self.conn.execute(insert_query, data_item)
```
The unpacking seems unnecessary since the two are joined immediately?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1365275358
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,32 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(self, table: str, input_data: Iterable[tuple[str, list[float]]], vector_size: int) -> None:
+ """
+ Store embedding vector in Postgres table.
+
+ :param table: The Name of the table
+ :param input_data: Iterable containing tuples of input data and corresponding embedding vectors.
+ :param vector_size: The size of vector. The maximum dimensions can be 2,000
+ """
+ from pgvector.psycopg import register_vector
+ from psycopg2 import sql
+
+ self.conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
Review Comment:
```suggestion
```
AFAIK, before Postgres 13 CREATE EXTENSION required superuser permissions, in PG 13+ trusted extension could be installed by someone with appropriate CREATE, however it is not a case of [pgvector](https://github.com/pgvector/pgvector/issues/288)
So let keep create extensions for DBAs or for some one with appropriate permissions. We should not force users to use SU in their databases.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1356852725
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,29 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(
+ self, table: str, input_data: list[str], embeddings: list[float], vector_size: int
+ ) -> None:
+ """
+ Store embedding vector in Postgres table.
+
+ :param table: The Name of the table
+ :param input_data: The source data from which the embedding has been created
+ :param embeddings: The embedding vector response from LLM service
+ :param vector_size: The size of vector. The maximum dimensions can be 2,000
+ """
+ from pgvector.psycopg import register_vector
+
+ self.conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+ register_vector(self.conn)
+ self.conn.execute(
+ "CREATE TABLE IF NOT EXISTS %s (id bigserial PRIMARY KEY, content text, embedding vector(%s))",
+ (table, vector_size),
+ )
+
+ for content, embedding in zip(input_data, embeddings):
+ self.execute("INSERT INTO %s (content, embedding) VALUES (%s, %s)", (table, content, embedding))
Review Comment:
https://www.psycopg.org/docs/sql.html#module-psycopg2.sql
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1357810354
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,36 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(
+ self, table: str, input_data: list[str], embeddings: list[float], vector_size: int
Review Comment:
Instead of accepting two lists I feel it’d be more natural to accept a list of two tuples. Or better, an iterable since that’s the only required interface (for the for loop).
i.e.
```suggestion
self,
table: str,
input_data: Iterable[tuple[str, float]],
vector_size: int,
```
##########
airflow/providers/postgres/hooks/postgres.py:
##########
@@ -320,6 +320,36 @@ def _generate_insert_sql(
return sql
+ def ingest_embedding(
+ self, table: str, input_data: list[str], embeddings: list[float], vector_size: int
Review Comment:
Instead of accepting two lists I feel it’d be more natural to accept a list of two tuples. Or better, an iterable since that’s the only required interface (for the for loop).
i.e.
```suggestion
self,
table: str,
input_data: Iterable[tuple[str, float]],
vector_size: int,
```
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1356844234
##########
airflow/providers/postgres/provider.yaml:
##########
@@ -53,6 +53,7 @@ dependencies:
- apache-airflow>=2.4.0
- apache-airflow-providers-common-sql>=1.3.1
- psycopg2-binary>=2.8.0
+ - pgvector
Review Comment:
probably should be extra dependency rather than a dependncy?
##########
airflow/providers/postgres/provider.yaml:
##########
@@ -53,6 +53,7 @@ dependencies:
- apache-airflow>=2.4.0
- apache-airflow-providers-common-sql>=1.3.1
- psycopg2-binary>=2.8.0
+ - pgvector
Review Comment:
probably should be extra dependency rather than a dependancy?
##########
airflow/providers/postgres/provider.yaml:
##########
@@ -53,6 +53,7 @@ dependencies:
- apache-airflow>=2.4.0
- apache-airflow-providers-common-sql>=1.3.1
- psycopg2-binary>=2.8.0
+ - pgvector
Review Comment:
probably should be extra dependency rather than a dependency?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "mpgreg (via GitHub)" <gi...@apache.org>.
mpgreg commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1365231569
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -80,3 +86,60 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
+ """
+ Operator for ingesting text and embeddings into a PostgreSQL database using the pgvector library.
+
+ :param conn_id: The connection ID for the postgresql database.
+ :param input_data: Tuple containing the string input content and corresponding list of float vector
+ embeddings.
+ :param input_callable: A callable that returns a tuple containing the string input content and
+ corresponding list of float vector embeddings, if ``input_data`` is not provided.
+ :param input_callable_args: Positional arguments for the 'input_callable'.
+ :param input_callable_kwargs: Keyword arguments for the 'input_callable'.
+ :param kwargs: Additional keyword arguments for the BaseOperator.
+ """
+
+ def __init__(
+ self,
+ conn_id: str,
+ input_data: tuple[str, list[float]] | None = None,
+ input_callable: Callable[[Any], Any] | None = None,
+ input_callable_args: Collection[Any] | None = None,
+ input_callable_kwargs: Mapping[str, Any] | None = None,
Review Comment:
Why not just have an `@task.pgvector_import()` decorator which provides a cleaner UX for instantiating the hook and passing params? Conceptually like https://github.com/astronomer/ask-astro/blob/00cfbd7a48aafe5603b1ef49342f1dd68c148156/airflow/dags/ingestion/ask-astro-load-blogs.py#L167
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable pgvector support for Postgres provider [airflow]
Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis commented on code in PR #34891:
URL: https://github.com/apache/airflow/pull/34891#discussion_r1362910782
##########
airflow/providers/postgres/operators/postgres.py:
##########
@@ -80,3 +86,60 @@ def __init__(
AirflowProviderDeprecationWarning,
stacklevel=2,
)
+
+
+class PgVectorIngestOperator(BaseOperator):
+ """
+ Operator for ingesting text and embeddings into a PostgreSQL database using the pgvector library.
+
+ :param conn_id: The connection ID for the postgresql database.
+ :param input_data: Tuple containing the string input content and corresponding list of float vector
+ embeddings.
+ :param input_callable: A callable that returns a tuple containing the string input content and
+ corresponding list of float vector embeddings, if ``input_data`` is not provided.
+ :param input_callable_args: Positional arguments for the 'input_callable'.
+ :param input_callable_kwargs: Keyword arguments for the 'input_callable'.
+ :param kwargs: Additional keyword arguments for the BaseOperator.
+ """
+
+ def __init__(
+ self,
+ conn_id: str,
+ input_data: tuple[str, list[float]] | None = None,
+ input_callable: Callable[[Any], Any] | None = None,
+ input_callable_args: Collection[Any] | None = None,
+ input_callable_kwargs: Mapping[str, Any] | None = None,
Review Comment:
I wan't to address the same question as I asked in https://github.com/apache/airflow/pull/34921#discussion_r1358525838
What benefits provide this arguments, because for me it could be replaced be replaced by different things
1. Make input data templated field and provide thought XCom from upstream task.
2. Combination of taskflow (or PythonOperator) + PostgresHook
```python
@task
def awesome_task(conn_id: str):
hook = PostgresHook(postgres_conn_id=conn_id)
...
hook.ingest_embedding(
table="foo",
input_data=...
vector_size=42
)
```
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org