shahar1 commented on code in PR #50750:
URL: https://github.com/apache/airflow/pull/50750#discussion_r2094599220
##########
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:
Don't we want to raise an exception in this case?
If not, I think that logging fits better in this case:
```suggestion
self.log.error("Number of failed imports: %s",
len(failed_references))
```
--
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]
