sjyangkevin commented on code in PR #50750:
URL: https://github.com/apache/airflow/pull/50750#discussion_r2094602774
##########
providers/weaviate/src/airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -268,6 +269,61 @@ def _convert_dataframe_to_list(data: list[dict[str, Any]]
| pd.DataFrame | None)
data = json.loads(data.to_json(orient="records"))
return cast("list[dict[str, Any]]", data)
+ def batch_create_links(
+ self,
+ collection_name: str,
+ data: list[dict[str, Any]] | pd.DataFrame | None,
+ from_property_col: str = "from_property",
+ from_uuid_col: str = "from_uuid",
+ to_uuid_col: str = "to",
+ retry_attempts_per_object: int = 5,
+ ) -> list[ErrorReference] | None:
+ """
+ Batch create links from an object to another other object through
cross-references
(https://weaviate.io/developers/weaviate/manage-data/import#import-with-references).
+
+ :param collection_name: The name of the collection containing the
source objects.
+ :param data: list or dataframe of objects we want to create links.
+ :param from_property_col: name of the reference property column.
+ :param from_uuid_col: Name of the column containing the from UUID.
+ :param to_uuid_col: Name of the column containing the target UUID.
+ :param retry_attempts_per_object: number of time to try in case of
failure before giving up.
+ """
+ converted_data = self._convert_dataframe_to_list(data)
+ collection = self.get_collection(collection_name)
+
+ with collection.batch.dynamic() as batch:
+ # Batch create links
+ for data_obj in converted_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:
+ from_property = data_obj.pop(from_property_col, None)
+ from_uuid = data_obj.pop(from_uuid_col, None)
+ to_uuid = data_obj.pop(to_uuid_col, None)
+ self.log.debug(
+ "Attempt %s of create links between %s and %s
using reference property %s",
+ attempt.retry_state.attempt_number,
+ from_uuid,
+ to_uuid,
+ from_property,
+ )
+ batch.add_reference(
+ from_property=from_property,
+ from_uuid=from_uuid,
+ to=to_uuid,
+ )
+
+ failed_references = collection.batch.failed_references
+ if failed_references:
+ print(f"Number of failed imports: {len(failed_references)}")
Review Comment:
Thank you for pointing out. I agree that logging fits better.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
