sunank200 commented on code in PR #36177:
URL: https://github.com/apache/airflow/pull/36177#discussion_r1424884362
##########
airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -396,28 +399,46 @@ def batch_data(
.. seealso:: `batch_config_params options
<https://weaviate-python-client.readthedocs.io/en/v3.25.3/weaviate.batch.html#weaviate.batch.Batch.configure>`__
:param vector_col: name of the column containing the vector.
:param retry_attempts_per_object: number of time to try in case of
failure before giving up.
+ :param tenant: The tenant to which the object will be added.
+ :param uuid_col: Name of the column containing the UUID.
"""
client = self.conn
if not batch_config_params:
batch_config_params = {}
client.batch.configure(**batch_config_params)
data = self._convert_dataframe_to_list(data)
+ insertion_errors = []
with client.batch as batch:
# Batch import all data
- for index, data_obj in enumerate(data):
- for attempt in Retrying(
- stop=stop_after_attempt(retry_attempts_per_object),
- retry=(
- retry_if_exception(lambda exc:
check_http_error_is_retryable(exc))
- | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES)
- ),
- ):
- with attempt:
- self.log.debug(
- "Attempt %s of importing data: %s",
attempt.retry_state.attempt_number, index + 1
- )
- vector = data_obj.pop(vector_col, None)
- batch.add_data_object(data_obj, class_name,
vector=vector)
+ try:
Review Comment:
+1 tests for all error scenarios is required here.
##########
airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -396,28 +399,46 @@ def batch_data(
.. seealso:: `batch_config_params options
<https://weaviate-python-client.readthedocs.io/en/v3.25.3/weaviate.batch.html#weaviate.batch.Batch.configure>`__
:param vector_col: name of the column containing the vector.
:param retry_attempts_per_object: number of time to try in case of
failure before giving up.
+ :param tenant: The tenant to which the object will be added.
Review Comment:
For ingesting data, it should determine objects to insert, objects to delete
and objects that remain unchanged for each chunk. This would simplify the logic.
##########
airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -606,3 +627,220 @@ def object_exists(self, uuid: str | UUID, **kwargs) ->
bool:
"""
client = self.conn
return client.data_object.exists(uuid, **kwargs)
+
+ def _generate_uuids(
+ self,
+ df: pd.DataFrame,
+ class_name: str,
+ unique_columns: list[str],
+ vector_column: str | None = None,
+ uuid_column: str | None = None,
+ ) -> tuple[pd.DataFrame, str]:
+ """
+ Adds UUIDs to a DataFrame, useful for replace operations where UUIDs
must be known before ingestion.
+
+ By default, UUIDs are generated using a custom function if
'uuid_column' is not specified.
+ The function can potentially ingest the same data multiple times with
different UUIDs.
+
+ :param df: A dataframe with data to generate a UUID from.
+ :param class_name: The name of the class use as part of the uuid
namespace.
+ :param uuid_column: Name of the column to create. Default is 'id'.
+ :param unique_columns: A list of columns to use for UUID generation.
By default, all columns except
+ vector_column will be used.
+ :param vector_column: Name of the column containing the vector data.
If specified the vector will be
+ removed prior to generating the uuid.
+ """
+ column_names = df.columns.to_list()
+
+ difference_columns =
set(unique_columns).difference(set(df.columns.to_list()))
+ if difference_columns:
+ raise ValueError(f"Columns {', '.join(difference_columns)} don't
exist in dataframe")
+
+ if uuid_column is None:
+ self.log.info("No uuid_column provided. Generating UUIDs as column
name `id`.")
+ if "id" in column_names:
+ raise ValueError(
+ "Property 'id' already in dataset. Consider renaming or
specify 'uuid_column'."
+ )
+ else:
+ uuid_column = "id"
+
+ if uuid_column in column_names:
+ raise ValueError(
+ f"Property {uuid_column} already in dataset. Consider renaming
or specify a different"
+ f" 'uuid_column'."
+ )
+
+ df[uuid_column] = (
+ df[unique_columns]
+ .drop(columns=[vector_column], inplace=False, errors="ignore")
+ .apply(lambda row: generate_uuid5(identifier=row.to_dict(),
namespace=class_name), axis=1)
+ )
+
+ return df, uuid_column
+
+ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str,
class_name: str, existing: str):
+ """
+ Helper function to check if the objects with uuid exist or not.
+
+ :param data: A single pandas DataFrame.
+ :param uuid_column: Column with pre-generated UUIDs.
+ :param class_name: Name of the class in Weaviate schema where data is
to be ingested.
+ :param existing: Strategy for handling existing data: 'skip', or
'replace'.
+ """
+ existing_uuid: set = set()
+ non_existing_uuid: set = set()
+
+ if existing == "replace":
+ existing_uuid = set(data[uuid_column].to_list())
+ else:
+ self.log.info("checking if %s objects exists.", data.shape[0])
+ for uuid in data[uuid_column]:
+ for attempt in Retrying(
+ stop=stop_after_attempt(5),
+ retry=(
+ retry_if_exception(lambda exc:
check_http_error_is_retryable(exc))
+ | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES)
+ ),
+ ):
+ with attempt:
+ if self.object_exists(uuid=uuid,
class_name=class_name):
+ existing_uuid.add(uuid)
+ if existing == "error":
+ return existing_uuid, non_existing_uuid
+ self.log.debug("object with uuid %s exists.", uuid)
+ else:
+ non_existing_uuid.add(uuid)
+ self.log.debug("object with uuid %s don't
exists.", uuid)
+
+ self.log.info(
+ f"Objects to override {len(existing_uuid)} and
{len(non_existing_uuid)} " f"objects to create"
+ )
+ return existing_uuid, non_existing_uuid
+
+ def _delete_objects(self, uuids: Collection, class_name: str,
retry_attempts_per_object: int = 5):
+ """
+ Helper function for `create_or_replace_objects()` to delete multiple
objects.
+
+ :param uuids: Collection of uuids.
+ :param class_name: Name of the class in Weaviate schema where data is
to be ingested.
+ :param retry_attempts_per_object: number of time to try in case of
failure before giving up.
+ """
+ for uuid in uuids:
+ for attempt in Retrying(
+ stop=stop_after_attempt(retry_attempts_per_object),
+ retry=(
+ retry_if_exception(lambda exc:
check_http_error_is_retryable(exc))
+ | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES)
+ ),
+ ):
+ with attempt:
+ try:
+ self.delete_object(uuid=uuid, class_name=class_name)
+ self.log.debug("Deleted object with uuid %s", uuid)
+ except weaviate.exceptions.UnexpectedStatusCodeException
as e:
+ if e.status_code == 404:
+ self.log.debug("Tried to delete a non existent
object with uuid %s", uuid)
+ else:
+ self.log.debug("Error occurred while trying to
delete object with uuid %s", uuid)
+ raise e
+
+ self.log.info("Deleted %s objects.", len(uuids))
+
+ def create_or_replace_objects(
+ self,
+ data: pd.DataFrame | list[dict[str, Any]],
+ class_name: str,
+ existing: str = "skip",
+ unique_columns: list[str] | str | None = None,
+ uuid_column: str | None = None,
+ vector_column: str = "Vector",
+ batch_config_params: dict | None = None,
+ tenant: str | None = None,
+ ) -> list:
+ """
+ create or replace objects.
+
+ Provides users with multiple ways of dealing with existing values.
+ 1. replace: replace the existing object with new object. This
option requires to identify the rows
+ uniquely, which by default is done by using all columns(except
`vector`) to create a uuid.
+ User can modify this behaviour by providing `unique_columns`
params.
+ 2. skip: skip the existing objects.
+ 3. error: raise an error if an existing object is found.
+
+
+ :param data: A single pandas DataFrame or a list of dicts to be
ingested.
+ :param class_name: Name of the class in Weaviate schema where data is
to be ingested.
+ :param existing: Strategy for handling existing data: 'skip', or
'replace'. Default is 'skip'.
+ :param unique_columns: Columns in DataFrame or keys in dict uniquely
identifying each document,
+ required for 'upsert' operations.
+ :param uuid_column: Column with pre-generated UUIDs. If not provided,
UUIDs will be generated.
+ :param vector_column: Column with embedding vectors for pre-embedded
data.
+ :param batch_config_params: Additional parameters for Weaviate batch
configuration.
+ :param tenant: The tenant to which the object will be added.
+ :return: list of UUID which failed to create
+ """
+ import pandas as pd
+
+ if existing not in ["skip", "replace", "error"]:
+ raise ValueError("Invalid parameter for 'existing'. Choices are
'skip', 'replace', 'error'.")
+
+ if isinstance(data, list):
+ data = pd.json_normalize(data)
+
+ if isinstance(unique_columns, str):
+ unique_columns = [unique_columns]
+ elif unique_columns is None:
+ unique_columns = sorted(data.columns.to_list())
+
+ self.log.info("Inserting %s objects.", data.shape[0])
+
+ if uuid_column is None or uuid_column not in data.columns:
+ (
+ data,
+ uuid_column,
+ ) = self._generate_uuids(
+ df=data,
+ class_name=class_name,
+ unique_columns=unique_columns,
+ vector_column=vector_column,
+ uuid_column=uuid_column,
+ )
+ # drop duplicate rows, using uuid_column and unique_columns
+ data = data.drop_duplicates(
+ subset=list({*unique_columns, uuid_column} - {vector_column,
None}), keep="first"
+ )
+
+ uuids_to_create = set()
+ existing_uuid, non_existing_uuid = self._check_existing_objects(
+ data=data, uuid_column=uuid_column, class_name=class_name,
existing=existing
+ )
+ if existing == "error" and len(existing_uuid):
+ self.log.info("Found duplicate UUIDs %s", " ,".join(existing_uuid))
+ raise ValueError(
+ f"Found {len(existing_uuid)} object with duplicate UUIDs. You
can either ignore or replace"
+ f" them by passing 'existing=skip' or 'existing=replace'
respectively."
+ )
+ elif existing == "replace":
+ uuids_to_create = existing_uuid.union(non_existing_uuid)
+ self._delete_objects(existing_uuid, class_name=class_name)
Review Comment:
+1
##########
airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -606,3 +627,220 @@ def object_exists(self, uuid: str | UUID, **kwargs) ->
bool:
"""
client = self.conn
return client.data_object.exists(uuid, **kwargs)
+
+ def _generate_uuids(
+ self,
+ df: pd.DataFrame,
+ class_name: str,
+ unique_columns: list[str],
+ vector_column: str | None = None,
+ uuid_column: str | None = None,
+ ) -> tuple[pd.DataFrame, str]:
+ """
+ Adds UUIDs to a DataFrame, useful for replace operations where UUIDs
must be known before ingestion.
+
+ By default, UUIDs are generated using a custom function if
'uuid_column' is not specified.
Review Comment:
We use
[generate_uuid](https://github.com/astronomer/ask-astro/blob/main/airflow/include/tasks/extract/utils/weaviate/ask_astro_weaviate_hook.py#L162)
in Ask-Astro as well for the reference
##########
airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -396,28 +399,46 @@ def batch_data(
.. seealso:: `batch_config_params options
<https://weaviate-python-client.readthedocs.io/en/v3.25.3/weaviate.batch.html#weaviate.batch.Batch.configure>`__
:param vector_col: name of the column containing the vector.
:param retry_attempts_per_object: number of time to try in case of
failure before giving up.
+ :param tenant: The tenant to which the object will be added.
+ :param uuid_col: Name of the column containing the UUID.
"""
client = self.conn
if not batch_config_params:
batch_config_params = {}
client.batch.configure(**batch_config_params)
data = self._convert_dataframe_to_list(data)
+ insertion_errors = []
with client.batch as batch:
# Batch import all data
- for index, data_obj in enumerate(data):
- for attempt in Retrying(
- stop=stop_after_attempt(retry_attempts_per_object),
- retry=(
- retry_if_exception(lambda exc:
check_http_error_is_retryable(exc))
- | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES)
- ),
- ):
- with attempt:
- self.log.debug(
- "Attempt %s of importing data: %s",
attempt.retry_state.attempt_number, index + 1
- )
- vector = data_obj.pop(vector_col, None)
- batch.add_data_object(data_obj, class_name,
vector=vector)
+ try:
+ for index, data_obj in enumerate(data):
+ for attempt in Retrying(
+ stop=stop_after_attempt(retry_attempts_per_object),
+ retry=(
+ retry_if_exception(lambda exc:
check_http_error_is_retryable(exc))
+ |
retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES)
+ ),
+ ):
+ with attempt:
+ vector = data_obj.pop(vector_col, None)
+ uuid = data_obj.pop(uuid_col, None)
+ self.log.debug(
+ "Attempt %s of inserting object with uuid: %s",
+ attempt.retry_state.attempt_number,
+ uuid,
+ )
+ batch.add_data_object(
+ data_object=data_obj,
+ class_name=class_name,
+ vector=vector,
+ uuid=uuid,
+ tenant=tenant,
+ )
+ self.log.debug("Inserted object with uuid: %s",
uuid)
Review Comment:
Yes the insert is completed only after context manager is over.
##########
airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -606,3 +627,220 @@ def object_exists(self, uuid: str | UUID, **kwargs) ->
bool:
"""
client = self.conn
return client.data_object.exists(uuid, **kwargs)
+
+ def _generate_uuids(
+ self,
+ df: pd.DataFrame,
+ class_name: str,
+ unique_columns: list[str],
+ vector_column: str | None = None,
+ uuid_column: str | None = None,
+ ) -> tuple[pd.DataFrame, str]:
+ """
+ Adds UUIDs to a DataFrame, useful for replace operations where UUIDs
must be known before ingestion.
+
+ By default, UUIDs are generated using a custom function if
'uuid_column' is not specified.
+ The function can potentially ingest the same data multiple times with
different UUIDs.
+
+ :param df: A dataframe with data to generate a UUID from.
+ :param class_name: The name of the class use as part of the uuid
namespace.
+ :param uuid_column: Name of the column to create. Default is 'id'.
+ :param unique_columns: A list of columns to use for UUID generation.
By default, all columns except
+ vector_column will be used.
+ :param vector_column: Name of the column containing the vector data.
If specified the vector will be
+ removed prior to generating the uuid.
+ """
+ column_names = df.columns.to_list()
+
+ difference_columns =
set(unique_columns).difference(set(df.columns.to_list()))
+ if difference_columns:
+ raise ValueError(f"Columns {', '.join(difference_columns)} don't
exist in dataframe")
+
+ if uuid_column is None:
+ self.log.info("No uuid_column provided. Generating UUIDs as column
name `id`.")
+ if "id" in column_names:
+ raise ValueError(
+ "Property 'id' already in dataset. Consider renaming or
specify 'uuid_column'."
+ )
+ else:
+ uuid_column = "id"
+
+ if uuid_column in column_names:
+ raise ValueError(
+ f"Property {uuid_column} already in dataset. Consider renaming
or specify a different"
+ f" 'uuid_column'."
+ )
+
+ df[uuid_column] = (
+ df[unique_columns]
+ .drop(columns=[vector_column], inplace=False, errors="ignore")
+ .apply(lambda row: generate_uuid5(identifier=row.to_dict(),
namespace=class_name), axis=1)
+ )
+
+ return df, uuid_column
+
+ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str,
class_name: str, existing: str):
Review Comment:
The upsert is a must here. It should have strategies to "upsert", "skip",
"replace" and "error". The scenarios mentioned by @mpgreg are must for
ingestion for any real use-case for LLM applications
--
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
